Uncertainty-Aware Control
for Reliable AI Systems

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional, NamedTuple
from enum import Enum

class Decision(Enum):
    PREDICT = "predict"
    ABSTAIN = "abstain"
    DEFER = "defer"

class UncertaintyEstimate(NamedTuple):
    aleatoric: torch.Tensor
    epistemic: torch.Tensor
    total: torch.Tensor


class UncertaintyAwareController(nn.Module):
    """Unified uncertainty quantification and control system."""
    
    def __init__(
        self,
        base_model: nn.Module,
        num_ensemble: int = 5,
        abstain_threshold: float = 0.85,
        defer_cost: float = 0.1,
        num_mc_samples: int = 20
    ):
        super().__init__()
        self.ensemble = nn.ModuleList([
            self._clone_model(base_model) for _ in range(num_ensemble)
        ])
        self.abstain_threshold = abstain_threshold
        self.defer_cost = defer_cost
        self.num_mc_samples = num_mc_samples
        
        # Calibration network
        self.calibrator = nn.Sequential(
            nn.Linear(3, 32),  # conf, epistemic, aleatoric
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )
    
    def compute_uncertainty(
        self, 
        x: torch.Tensor
    ) -> Tuple[torch.Tensor, UncertaintyEstimate]:
        """Compute decomposed uncertainty from ensemble."""
        # Collect ensemble predictions
        all_probs = []
        for model in self.ensemble:
            model.eval()
            with torch.no_grad():
                logits = model(x)
                probs = F.softmax(logits, dim=-1)
                all_probs.append(probs)
        
        all_probs = torch.stack(all_probs)  # [M, B, C]
        
        # Mean prediction
        mean_probs = all_probs.mean(dim=0)  # [B, C]
        
        # Aleatoric: expected entropy
        entropies = -(all_probs * (all_probs + 1e-10).log()).sum(dim=-1)
        aleatoric = entropies.mean(dim=0)  # [B]
        
        # Total: entropy of mean
        total = -(mean_probs * (mean_probs + 1e-10).log()).sum(dim=-1)
        
        # Epistemic: mutual information
        epistemic = total - aleatoric
        
        uncertainty = UncertaintyEstimate(
            aleatoric=aleatoric,
            epistemic=epistemic,
            total=total
        )
        
        return mean_probs, uncertainty
    
    def calibrate_confidence(
        self,
        raw_conf: torch.Tensor,
        uncertainty: UncertaintyEstimate
    ) -> torch.Tensor:
        """Apply learned calibration to raw confidence."""
        features = torch.stack([
            raw_conf,
            uncertainty.epistemic,
            uncertainty.aleatoric
        ], dim=-1)
        
        return self.calibrator(features).squeeze(-1)
    
    def should_defer(
        self,
        epistemic: torch.Tensor,
        model_risk: torch.Tensor
    ) -> torch.Tensor:
        """Decide whether to defer to human expert."""
        # Estimate human would do better on high-epistemic cases
        human_risk = 0.05  # Assumed human error rate
        return (self.defer_cost + human_risk) < model_risk
    
    def forward(
        self,
        x: torch.Tensor
    ) -> Tuple[Decision, Optional[torch.Tensor], torch.Tensor]:
        """Make uncertainty-aware decision."""
        # Get predictions and uncertainty
        probs, uncertainty = self.compute_uncertainty(x)
        
        # Get predicted class and raw confidence
        raw_conf, pred = probs.max(dim=-1)
        
        # Calibrate confidence
        calibrated_conf = self.calibrate_confidence(raw_conf, uncertainty)
        
        # Estimate model risk
        model_risk = 1 - calibrated_conf
        
        # Decision logic
        decisions = []
        for i in range(x.size(0)):
            if calibrated_conf[i] >= self.abstain_threshold:
                decisions.append(Decision.PREDICT)
            elif self.should_defer(uncertainty.epistemic[i], model_risk[i]):
                decisions.append(Decision.DEFER)
            else:
                decisions.append(Decision.ABSTAIN)
        
        return decisions, pred, calibrated_conf


class SelectivePredictor(nn.Module):
    """Learn to predict and abstain jointly."""
    
    def __init__(self, base_model: nn.Module, num_classes: int):
        super().__init__()
        self.predictor = base_model
        self.selector = nn.Sequential(
            nn.Linear(num_classes + 128, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Returns predictions and selection scores."""
        features = self.predictor.get_features(x)
        logits = self.predictor.classifier(features)
        probs = F.softmax(logits, dim=-1)
        
        # Selection score: should we make a prediction?
        selector_input = torch.cat([probs, features], dim=-1)
        selection_score = self.selector(selector_input).squeeze(-1)
        
        return logits, selection_score
    
    def selective_loss(
        self,
        logits: torch.Tensor,
        selection: torch.Tensor,
        targets: torch.Tensor,
        coverage_target: float = 0.85,
        abstain_cost: float = 0.1
    ) -> torch.Tensor:
        """Compute selective prediction loss."""
        # Classification loss on selected samples
        ce_loss = F.cross_entropy(logits, targets, reduction='none')
        selective_ce = (ce_loss * selection).sum() / (selection.sum() + 1e-10)
        
        # Abstention penalty
        abstain_penalty = abstain_cost * (1 - selection).mean()
        
        # Coverage constraint
        coverage = selection.mean()
        coverage_penalty = F.relu(coverage_target - coverage) ** 2
        
        return selective_ce + abstain_penalty + 10 * coverage_penalty