how to decide which decision should be taken from four given options?
For my snake AI, I'm using pytorch whose structure is:

class SnakeAI(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(24, 16, bias=True),
            nn.ReLU(),
            nn.Linear(16, 8, bias=True),
            nn.ReLU(),
            nn.Linear(8, 4, bias=True),
            nn.Sigmoid()
        )

    def forward(self, inputs):
        x = self.fc(inputs)
        return x


Based on the inputs it gives me 4 values which to my understanding should be the probability of which decision to be, whose sum should be equal to 1. which is not the case here as the output is:

[0.59388083 0.5833764  0.47855872 0.5388371 ]


based on the given output should i use argmax to decide which direction to go with? is there any better way to do that?

You should change your last layer to nn.Softmax(-1). Then you would get in return a probability distribution and apply torch.max() which returns the maximum value and its index
