1"""Losses: next-token cross-entropy (per position), controller regression/classification losses."""2from __future__ import annotations34import torch5import torch.nn.functional as F678def token_nll(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = -100) -> torch.Tensor:9 """Per-position negative log-likelihood [B,T] (float32); ignored positions get 0."""10 B, T, V = logits.shape11 nll = F.cross_entropy(logits.float().reshape(B * T, V), targets.reshape(B * T), ignore_index=ignore_index, reduction="none")12 return nll.view(B, T)131415def lm_loss(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = -100) -> torch.Tensor:16 nll = token_nll(logits, targets, ignore_index)17 mask = (targets != ignore_index).float()18 return (nll * mask).sum() / mask.sum().clamp(min=1.0)192021def controller_losses(preds: dict, g_obs: torch.Tensor, valid: torch.Tensor, huber_delta: float = 0.1):22 """preds: dict of [N] tensors for the logged (state, action) pairs; g_obs: [N] observed gains (detached);23 valid: [N] bool. Returns dict(gain=Huber, harm=BCE, nll=Gaussian NLL for sigma)."""24 g_obs = g_obs.detach()25 v = valid.float()26 n = v.sum().clamp(min=1.0)27 gain = (F.huber_loss(preds["g_hat"], g_obs, delta=huber_delta, reduction="none") * v).sum() / n28 harm_t = (g_obs < 0).float()29 harm = (F.binary_cross_entropy_with_logits(preds["harm_logit"], harm_t, reduction="none") * v).sum() / n30 log_sigma = preds["log_sigma"].clamp(-7, 3)31 nll = ((0.5 * ((g_obs - preds["g_hat"].detach()) ** 2) * torch.exp(-2 * log_sigma) + log_sigma) * v).sum() / n32 return {"gain": gain, "harm": harm, "nll": nll}33