model.py

import torch
import torch.nn as nn

class UncertaintyActionModel(nn.Module):
    def init(self):
        super().init()
        self.fc = nn.Linear(3, 3)  # confidence, risk, ambiguity

    def forward(self, x):
        return torch.softmax(self.fc(x), dim=-1)