Update model.py - fix heteroscedastic loss clamping

model.py

@@ -727,11 +727,10 @@ class NextStateLoss(nn.Module):
         # --- Heteroscedastic regularization (learned aleatoric uncertainty) ---
         if 'log_var' in predictions:
             log_var = predictions['log_var'][:, :-1, :]  # (B, L-1, 6)
-            # Regularize: penalize overly high uncertainty (prevent collapse)
-            # The individual heads already implicitly learn to attend to uncertainty
-            # via the gradient signal, but we add a mild KL-like penalty
-            log_var_penalty = 0.01 * log_var.mean()
-            losses['log_var_reg'] = log_var_penalty
+            # Clamp log_var to prevent collapse: [-5, 5] range
+            log_var_clamped = torch.clamp(log_var, -5.0, 5.0)
+            # Regularize toward 0 (unit variance prior)
+            losses['log_var_reg'] = 0.1 * (log_var_clamped ** 2).mean()
         
         # Total loss
         total_loss = sum(losses.values())