Using toch log_prob to calculate the probability of selecting multiple values of the distribution itself
I try to use log_prob to get the probability of selecting a value from a normal distribution,
I got dist from a neural network and action from dist.sample()
In a learning phase, I give 5 tensors to a neural network, and it gives me 5 dist, and from dists, I got 5 actions. The problem is that I want to select an action over its own distribution, but this function gives me the probability of action in all distributions. The data on the diameter of the output matrix is the values I want, but I wonder if there is an easy way to implement this part?
I use this block of code:
states = T.tensor(state[b], dtype=T.float).to(agent.device)
old_probs = T.tensor(log_prob[b]).to(agent.device)
actions = T.tensor(action[b]).to(agent.device)
values = T.tensor(value[b]).to(agent.device)

dist = actor(states)
new_probs = dist.log_prob(actions)

and the output is
tensor([[-1.1823, -0.9680, -3.6280, -1.1112, -1.9610],
        [-1.5279, -1.1463, -2.5806, -1.0561, -1.4768],
        [-1.6258, -1.1618, -2.5027, -1.0100, -1.3882],
        [-1.6125, -1.1576, -2.5169, -1.0133, -1.3989],
        [-1.3384, -1.0965, -2.9404, -1.1370, -1.7129]], device='cuda:0',
       dtype=torch.float64, grad_fn=&lt;SubBackward0&gt;)

but the output must be like:
tensor([-1.1823, -1.1463, -2.5027, -1.0133, -1.7129], device='cuda:0',
       grad_fn=&lt;SqueezeBackward1&gt;)


You can select the diagonal of your matrix with torch.diag:
&gt;&gt;&gt; new_probs.diag()
tensor([-1.1823, -1.1463, -2.5027, -1.0133, -1.7129], 
  device='cuda:0', grad_fn=&lt;DiagBackward&gt;)

