import torch


def transform_u(logits, driver_matrix, budget):
    """
    A function that transforms RL logits to valid controls.
    :param logits: RL action logits
    :param driver_matrix: driver matrix for control selection
    :budget: total budget available
    """
    logits = torch.matmul(driver_matrix.unsqueeze(0), logits.unsqueeze(-1)).squeeze(-1)
    u = torch.nn.functional.softmax(logits, dim=-1) \
        * budget \
        * driver_matrix.sum(-1)
    return u