I am trying to implement a basic Policy Gradient training setup. I saw some examples for games that had always the same possible actions in each state. However I am wondering how one would implement the action sampling from the policy when in some states not all actions are allowed.
Right now I do the action sampling like this:
action = torch.multinomial(policy, 1)
log_probs = torch.log(torch.gather(policy, -1, action))

For the limited action set I thought of something like this
limited_policy = policy[valids]
action = torch.multinomial(limited_policy, 1)
log_probs = torch.log(torch.gather(limited_policy, -1, action))

Here the problem is that the shape of the limited policy changes through the valid indexing and therefore not all rows have the same length anymore and then I cannot use it as a batch anymore.
So what is the right way of handling the limited action space that it also handles the gradients properly?
EDIT
I now set the probability actions manually to zero for not allowed actions like this:
policy[1 - valids] = 0
action = torch.multinomial(policy, 1)
log_probs = torch.log(torch.gather(policy, -1, action))

But I would prefer if there is a solution that somehow consideres this in the sampling function. So that the network can actually “learn” that these actions are not allowed or less wanted in some states.