loss

loss

Loss configuration for policy training.

BaseLossConfig dataclass

BaseLossConfig(_target_=MISSING)

Base configuration for loss modules.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.

RegressionLossConfig dataclass

RegressionLossConfig(_target_='versatil.metrics.losses.regression.RegressionLoss', action_keys=MISSING, mse_weight=1.0, l1_weight=0.0, huber_weight=0.0, huber_delta=1.0, per_key_weights=None)

Bases: BaseLossConfig

Configuration for regression loss (position, orientation).

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
action_keys	list[str]	List of action keys to compute loss for (e.g., ['position', 'orientation']).
mse_weight	float	Weight for MSE loss.
l1_weight	float	Weight for L1 loss.
huber_weight	float	Weight for Huber loss.
huber_delta	float	Delta parameter for Huber loss.
per_key_weights	dict[str, float] | None	Optional dictionary of per-key weights.

GripperLossConfig dataclass

GripperLossConfig(_target_='versatil.metrics.losses.gripper.GripperLoss', key=MISSING, actions_metadata='${task.action_space.actions_metadata}', bce_weight=1.0, mse_weight=1.0, pos_weight=None)

Bases: BaseLossConfig

Configuration for gripper loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
key	str	Action key for gripper.
actions_metadata	Any	Dict of metadata of the action space.
bce_weight	float	Weight for binary cross entropy (binary gripper).
mse_weight	float	Weight for MSE loss (continuous gripper).
pos_weight	float | None	Optional positive class weight for BCE.

KLDivergenceLossConfig dataclass

KLDivergenceLossConfig(_target_='versatil.metrics.losses.divergence.KLDivergenceLoss', weight=0.0001, prior_regularization_weight=0.0)

Bases: BaseLossConfig

Configuration for KL divergence loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Weight for KL divergence loss KL(posterior || prior).
prior_regularization_weight	float	Weight for KL(prior || N(0,I)) regularization. Only meaningful for learned priors. Pushes the learned prior towards a standard Gaussian.

GaussianEntropyLossConfig dataclass

GaussianEntropyLossConfig(_target_='versatil.metrics.losses.divergence.GaussianEntropyLoss', key=MISSING, weight=0.0)

Bases: BaseLossConfig

Configuration for entropy loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
key	str	Prediction key for logvar tensor to compute entropy over.
weight	float	Loss weight. Positive values encourage higher entropy.

BinaryKLDivergenceLossConfig dataclass

BinaryKLDivergenceLossConfig(_target_='versatil.metrics.losses.divergence.BinaryKLDivergenceLoss', weight=0.0001, free_bits=0.0, latent_bits=MISSING, entropy_weight=0.005)

Bases: BaseLossConfig

Configuration for binary KL divergence loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Weight for KL divergence loss.
free_bits	float	Free bits threshold (only penalize KL above this value).
latent_bits	int	Number of bits of the latent codes.
entropy_weight	float	Weight for the entropy regularization term.

MaximumMeanDiscrepancyLossConfig dataclass

MaximumMeanDiscrepancyLossConfig(_target_='versatil.metrics.losses.maximum_mean_discrepancy.MaximumMeanDiscrepancyLoss', weight=1.0, prior_regularization_weight=0.0, prior_target_key='${latent_key:POSTERIOR_LATENT}', kernel_type=value, bandwidth_multipliers=(lambda: [0.2, 0.5, 1.0, 2.0, 5.0])(), use_median_heuristic=True, use_fixed_gaussian_as_prior=False)

Bases: BaseLossConfig

Configuration for Maximum Mean Discrepancy (MMD) loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Loss weight for MMD(posterior, prior).
prior_regularization_weight	float	Weight for MMD(prior, N(0,I)) regularization. Only meaningful for learned priors.
prior_target_key	str	Posterior output key used as aggregate prior-matching samples. Use LatentKey.POSTERIOR_MU for deterministic WAE-style matching.
kernel_type	str	Kernel type for MMD computation (see KernelType enum).
bandwidth_multipliers	list[float] | None	Scale factors for bandwidth. When use_median_heuristic=True these scale the adaptive median. When False these are absolute bandwidth values. WAE recommends [2 * latent_dim] with use_median_heuristic=False.
use_median_heuristic	bool	Adaptive bandwidth via median heuristic (True) or fixed absolute bandwidths (False).
use_fixed_gaussian_as_prior	bool	If True, always use standard Gaussian as prior.

ConditionalMaximumMeanDiscrepancyLossConfig dataclass

ConditionalMaximumMeanDiscrepancyLossConfig(_target_='versatil.metrics.losses.maximum_mean_discrepancy.ConditionalMaximumMeanDiscrepancyLoss', weight=1.0, state_weight=1.0, prior_target_key='${latent_key:POSTERIOR_LATENT}', condition_key='${latent_key:PRIOR_CONDITION}', kernel_type=value, bandwidth_multipliers=(lambda: [0.2, 0.5, 1.0, 2.0, 5.0])(), use_median_heuristic=True, condition_kernel_type=value, condition_bandwidth_multipliers=(lambda: [0.2, 0.5, 1.0, 2.0, 5.0])(), condition_use_median_heuristic=True, normalize_condition=True)

Bases: BaseLossConfig

Configuration for conditional state-latent MMD loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Scalar weight of this loss in the total loss.
state_weight	float	Weight of the state kernel in the joint kernel.
prior_target_key	str	Metadata key holding prior samples matched against the posterior.
condition_key	str	Feature key used as the conditioning variable.
kernel_type	str	Kernel for the latent term, rbf or imq.
bandwidth_multipliers	list[float] | None	Bandwidth multipliers of the latent kernel mixture.
use_median_heuristic	bool	Whether the latent bandwidth uses the median heuristic.
condition_kernel_type	str	Kernel for the conditioning term, rbf or imq.
condition_bandwidth_multipliers	list[float] | None	Bandwidth multipliers of the conditioning kernel mixture.
condition_use_median_heuristic	bool	Whether the conditioning bandwidth uses the median heuristic.
normalize_condition	bool	Whether the conditioning variable is standardized before the kernel.

VQCommitmentLossConfig dataclass

VQCommitmentLossConfig(_target_='versatil.metrics.losses.vector_quantization.VQCommitmentLoss', num_codes=MISSING, num_residual_layers=MISSING, weight=1.0)

Bases: BaseLossConfig

Configuration for VQ commitment loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
num_codes	int	Number of codebook entries per residual layer (K). Must match the VQ posterior's ResidualVQ configuration.
num_residual_layers	int	Number of residual VQ layers. Must match the VQ posterior's ResidualVQ configuration.
weight	float	Loss weight for the commitment term ||z_continuous - sg(z_quantized)||^2.

VQPriorCrossEntropyLossConfig dataclass

VQPriorCrossEntropyLossConfig(_target_='versatil.metrics.losses.vector_quantization.VQPriorCrossEntropyLoss', weight=1.0)

Bases: BaseLossConfig

Configuration for VQ prior cross-entropy loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Loss weight for the cross-entropy term.

BinaryMaximumMeanDiscrepancyLossConfig dataclass

BinaryMaximumMeanDiscrepancyLossConfig(_target_='versatil.metrics.losses.maximum_mean_discrepancy.BinaryMaximumMeanDiscrepancyLoss', weight=1.0)

Bases: BaseLossConfig

Configuration for Binary Maximum Mean Discrepancy (MMD) loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Loss weight.

TrajectoryLengthLossConfig dataclass

TrajectoryLengthLossConfig(_target_='versatil.metrics.losses.trajectory.TrajectoryLengthLoss', weight=0.1, action_key=MISSING)

Bases: BaseLossConfig

Configuration for trajectory length loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Weight for length loss.
action_key	str	Action key to compute length for.

TrajectorySmoothnessConfig dataclass

TrajectorySmoothnessConfig(_target_='versatil.metrics.losses.trajectory.TrajectorySmoothness', weight=0.01, action_key=MISSING)

Bases: BaseLossConfig

Configuration for trajectory smoothness loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Weight for smoothness loss.
action_key	str	Action key to compute smoothness for.

ActionTokenLossConfig dataclass

ActionTokenLossConfig(_target_='versatil.metrics.losses.classification.ActionTokenLoss', weight=1.0, label_smoothing=0.2)

Bases: BaseLossConfig

Configuration for action token cross-entropy loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Scalar multiplier applied to the cross-entropy term.
label_smoothing	float	Label smoothing factor [0, 1].

PhaseClassificationLossConfig dataclass

PhaseClassificationLossConfig(_target_='versatil.metrics.losses.classification.PhaseClassificationLoss', key=MISSING, cross_entropy_weight=1.0, entropy_weight=0.0, label_smoothing=0.0)

Bases: BaseLossConfig

Configuration for phase classification loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
key	str	Key for phase labels.
cross_entropy_weight	float	Weight for cross-entropy loss.
entropy_weight	float	Weight for entropy regularization (Entropy maximization avoids experts collapse).
label_smoothing	float	Label smoothing factor for cross-entropy.

GripperMixtureNLLossConfig dataclass

GripperMixtureNLLossConfig(_target_='versatil.metrics.losses.mixture.GripperMixtureNLLoss', key=MISSING, actions_metadata='${task.action_space.actions_metadata}', weight=1.0, learned_variance=False, sigma=0.5, min_variance=0.0001)

Bases: BaseLossConfig

Configuration for gripper Mixture Negative Log-Likelihood loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
key	str	Key for gripper actions.
actions_metadata	Any	Dict of metadata of the action space.
weight	float	Loss weight.
learned_variance	bool	If True, expects {key}_mean and {key}_logvar for continuous. If False, expects {key} (stacked means) and uses sigma.
sigma	float	Fixed std for continuous gripper (only used when learned_variance=False).
min_variance	float	Minimum variance for numerical stability (learned_variance=True).

CompositeLossConfig dataclass

CompositeLossConfig(_target_='versatil.metrics.losses.composite.CompositeLoss', loss_modules=dict(), weights=None)

Bases: BaseLossConfig

Configuration for composite loss with custom modules.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
loss_modules	dict[str, Any]	Dictionary of loss module names to loss instances.
weights	dict[str, float] | None	Deprecated legacy composite weights. Kept only for config compatibility and ignored at runtime.

PriorDenoisingLossConfig dataclass

PriorDenoisingLossConfig(_target_='versatil.metrics.losses.prior_denoising.PriorDenoisingLoss', weight=1.0)

Bases: BaseLossConfig

Configuration for diffusion prior denoising loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Weight for this loss component.

MoELossConfig dataclass

MoELossConfig(_target_='versatil.metrics.losses.mixture_of_experts.MoELoss', base_loss=MISSING, entropy_weight=0.0, load_balance_weight=0.0)

Configuration for Mixture of Experts (MoE) loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
base_loss	BaseLossConfig	Any BaseLoss instance to wrap (e.g., RegressionLoss(...)).
entropy_weight	float	Weight for per-example routing entropy. Penalizes peaky-per- example routing. Pushes each example's routing distribution toward uniform, which prevents one example from being routed to a single expert with probability 1.
load_balance_weight	float	Weight for Switch-Transformer-style load-balancing term. Penalizes batch-level imbalance in expert usage. The term is K * sum_k f_k * P_k where f_k is the fraction of examples whose argmax routes to expert k and P_k is the mean routing weight for expert k across the batch. Minimum value 1.0 is reached when usage is uniform across the batch. Crucially, this allows per-example routing to be peaky (so experts can specialize) while still forcing every expert to be used by some examples (so no expert dies). Use this when entropy alone produces dead experts.

GaussianMixtureNLLossConfig dataclass

GaussianMixtureNLLossConfig(_target_='versatil.metrics.losses.mixture.GaussianMixtureNLLoss', action_keys=MISSING, weight=1.0, per_key_weights=None, learned_variance=True, sigmas=None, min_variance=0.0001)

Bases: BaseLossConfig

Configuration for Gaussian Mixture Negative Log-Likelihood loss.

Attributes:

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

VICLatentLossConfig dataclass

VICLatentLossConfig(_target_='versatil.metrics.losses.latent_geometry.VICLatentLoss', key='${latent_key:POSTERIOR_MU}', covariance_weight=3.0, variance_weight=10.0, gamma=0.3)

Bases: BaseLossConfig

Configuration for VICReg-style covariance + variance loss.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
key	str	Prediction key for latent mu tensor.
covariance_weight	float	Weight for off-diagonal covariance penalty.
variance_weight	float	Weight for variance hinge loss.
gamma	float	Hinge threshold for per-dimension standard deviation.

PosteriorGeometryLossConfig dataclass

PosteriorGeometryLossConfig(_target_='versatil.metrics.losses.latent_geometry.PosteriorGeometryLoss', key='${latent_key:POSTERIOR_MU}', mean_weight=0.0, std_weight=0.0, target_std=1.0, max_std_weight=0.0, max_std=2.0, covariance_weight=0.0, epsilon=1e-06)

Bases: BaseLossConfig

Configuration for posterior latent moment regularization.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
key	str	Prediction key for latent vectors.
mean_weight	float	Weight for squared batch-mean penalty.
std_weight	float	Weight for squared deviation from target_std.
target_std	float	Desired per-dimension posterior standard deviation.
max_std_weight	float	Weight for hinge penalty above max_std.
max_std	float	Maximum tolerated per-dimension standard deviation.
covariance_weight	float	Weight for off-diagonal covariance penalty.
epsilon	float	Numerical epsilon for standard deviation.

OptimalTransportLossConfig dataclass

OptimalTransportLossConfig(_target_='versatil.metrics.losses.optimal_transport.OptimalTransportLoss', action_keys=MISSING, weight=1.0, p=2, blur_fraction=0.1, reach_multiplier=None, expected_std=1.0, time_scale=1.0)

Bases: BaseLossConfig

Configuration for Optimal Transport loss using Sinkhorn divergence.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
action_keys	list[str]	List of keys for action tensors in predictions and targets.
weight	float	Scaling factor for the total loss.
p	int	Exponent for the ground cost. p=1 gives ||a - a'||_2, p=2 gives (1/2) * ||a - a'||_2^2.
blur_fraction	float	Dimensionless Sinkhorn regularization, expressed as a fraction of the reference pairwise scale sqrt(2 * dim) * expected_std. GeomLoss recommends ~0.1.
reach_multiplier	float | None	Unbalanced OT scale, as a multiple of the reference pairwise scale. None keeps balanced OT. Typical values for mild outlier tolerance are 3.0-10.0.
expected_std	float	Expected per-dimension standard deviation of the action samples. For actions normalized to [-1, 1], use ~1/sqrt(3) ~ 0.577.
time_scale	float	Scaling factor for the linear time embedding concatenated to actions. time_scale=0 gives permutation-invariant OT over the horizon.

LatentOptimalTransportLossConfig dataclass

LatentOptimalTransportLossConfig(_target_='versatil.metrics.losses.optimal_transport.LatentOptimalTransportLoss', weight=1.0, prior_target_key='${latent_key:POSTERIOR_LATENT}', p=2, blur_fraction=0.1, reach_multiplier=None)

Bases: BaseLossConfig

Configuration for latent Sinkhorn divergence between posterior and prior.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Scaling factor for the total loss.
prior_target_key	str	Posterior output key used as aggregate prior-matching samples. Use LatentKey.POSTERIOR_MU for deterministic WAE-style matching.
p	int	Exponent for the ground cost. p=2 is standard for W_2-style regularization of latent distributions.
blur_fraction	float	Dimensionless Sinkhorn regularization, as a fraction of the reference pairwise scale sqrt(2 * dim).
reach_multiplier	float | None	Unbalanced OT scale, as a multiple of the reference pairwise scale. None keeps balanced OT.

RelaxedConditionalLatentOptimalTransportLossConfig dataclass

RelaxedConditionalLatentOptimalTransportLossConfig(_target_='versatil.metrics.losses.optimal_transport.RelaxedConditionalLatentOptimalTransportLoss', weight=1.0, prior_target_key='${latent_key:POSTERIOR_LATENT}', condition_key='${latent_key:PRIOR_CONDITION}', p=2, blur_fraction=0.1, reach_multiplier=None, state_weight=1.0, normalize_condition=True)

Bases: BaseLossConfig

Configuration for relaxed conditional latent Sinkhorn divergence.

Attributes:

Name	Type	Description
_target_	str	Import path instantiated by Hydra.
weight	float	Scalar weight of this loss in the total loss.
prior_target_key	str	Metadata key holding prior samples matched against the posterior.
condition_key	str	Feature key used as the conditioning variable.
p	int	Order of the transport cost.
blur_fraction	float	Sinkhorn blur as a fraction of the point-cloud diameter.
reach_multiplier	float | None	Unbalanced-transport reach as a multiple of the diameter.
state_weight	float	Weight of the state coordinates in the transport cost.
normalize_condition	bool	Whether the conditioning variable is standardized before transport.