Advantage computed the wrong way?

Question

Here is the code written by Maxim Lapan. I am reading his book (Deep Reinforcement Learning Hands-on). I have seen a line in his code which is really weird. In the accumulation of the policy gradient $$\partial \theta_{\pi} \gets \partial \theta_{\pi} + \nabla_{\theta}\log\pi_{\theta} (a_i | s_i) (R - V_{\theta}(s_i))$$ we have to compute the advantage $R - V_{\theta}(s_i)$. In line 138, maxim uses adv_v = vals_ref_v - value_v.detach(). Visually, it looks fine, but look at the shape of each term.

ipdb> adv_v.shape                                                                                                                            
torch.Size([128, 128])

ipdb> vals_ref_v.shape                                                                                                                       
torch.Size([128])

ipdb> values_v.detach().shape                                                                                                                
torch.Size([128, 1]) 

In a much simpler code, it is equivalent to

In [1]: import torch                                                            

In [2]: t1 = torch.tensor([1, 2, 3])                                            

In [3]: t2 = torch.tensor([[4], [5], [6]])                                      

In [4]: t1 - t2                                                                 
Out[4]: 
tensor([[-3, -2, -1],
        [-4, -3, -2],
        [-5, -4, -3]])

In [5]: t1 - t2.detach()                                                        
Out[5]: 
tensor([[-3, -2, -1],
        [-4, -3, -2],
        [-5, -4, -3]])

I have trained the agent with his code and it works perfectly fine. I am very confused why it is good practice and what it is doing. Could someone enlighten me on the line adv_v = vals_ref_v - value_v.detach()? For me, the right thing to do was adv_v = vals_ref_v - value_v.squeeze(-1).

Here is the full algorithm used in his book :

UPDATE

As you can see by the image, it is converging even though adv_v = vals_ref_v - value_v.detach() looks wrongly implemented. It is not done yet, but I will update the question later.

Answer 1

Yeah, it seems like it's a wrong implementation. vals_ref_v is a matrix of 1 row, and 128 columns. value_v.detach() is a matrix of 128 row

Answer 2

I changed the line adv_v = vals_ref_v - value_v.detach() to adv_v = vals_ref_v - value_v.squeeze(-1).detach(). It seems the convergence is much faster. According to the A2C algorithm, it is just logic to apply $Q(a, s) - V(s)$, where $Q(a, s)$ and $V(s)$ with the same shape.

The call to detach() is important here as we don't want to propagate the PG into our value approximation head.