protomotions.agents.common.supervision module#

Shared supervision losses for imitation and distillation.

class protomotions.agents.common.supervision.SupervisionLossType(value)[source]#

Bases: str, Enum

Supported supervised losses over configured prediction and target keys.

MSE = 'mse'#

DISCRETE_CROSS_ENTROPY = 'discrete_cross_entropy'#

DISCRETE_KL = 'discrete_kl'#

CONTINUOUS_GAUSSIAN_KL = 'continuous_gaussian_kl'#

class protomotions.agents.common.supervision.SupervisionLossConfig( loss_type=SupervisionLossType.MSE, prediction_key='privileged_action', target_key='expert_actions', prediction_logvar_key=None, target_logvar_key=None, label_smoothing=0.0, weight=1.0, log_prefix='supervision', enabled=True, extra=<factory>, )[source]#

Bases: object

Key-based supervised loss over model outputs and labels.

Distillation agents use prediction_key and target_key to select tensors from a TensorDict batch, so the same loss config can supervise actions, discrete latent tokens, or distribution parameters.

loss_type: SupervisionLossType = 'mse'#

prediction_key: str = 'privileged_action'#

target_key: str = 'expert_actions'#

prediction_logvar_key: str | None = None#

target_logvar_key: str | None = None#

label_smoothing: float = 0.0#

weight: float = 1.0#

log_prefix: str = 'supervision'#

enabled: bool = True#

extra: Dict[str, str]#

__init__( loss_type=SupervisionLossType.MSE, prediction_key='privileged_action', target_key='expert_actions', prediction_logvar_key=None, target_logvar_key=None, label_smoothing=0.0, weight=1.0, log_prefix='supervision', enabled=True, extra=<factory>, )#

protomotions.agents.common.supervision.compute_supervision_loss(batch, config)[source]#

Compute a configured supervised loss.