dongminlee94/deep_rl

(TRPO) Trivial KL divergence computation in hessian_vector_product.

Opened this issue · 0 comments

Thanks for your great implementation of TRPO.

This is helpful for those (including me) who want to implement TRPO by themselves because there are a few implementations.

I just found a small mistake in your implementation, specifically in the hessian_vector_product method in TRPO agent (please, see below).
I think the KL divergence computation is incorrect because it is computing the KL divergence between two same distributions.
When I printed out the variable kl, it was always zero.

As far as I understand, it is true that the exact Fisher information matrix (FIM) should be computed by using only the current parameters.
However, the trick in TRPO is to approximate the FIM by computing the Hessian of the KL divergence between the old and current policy. This is reasonable when the old and current parameters are close enough.
This can be found in Section 6 Practical Algorithm in the TRPO paper.

Can you take a look at this issue?

Best regards,
Dongjin Lee

def hessian_vector_product(self, obs, p, damping_coeff=0.1):
      p.detach()
      kl = self.gaussian_kl(old_policy=self.policy, new_policy=self.policy, obs=obs)
      kl_grad = torch.autograd.grad(kl, self.policy.parameters(), create_graph=True)
      kl_grad = self.flat_grad(kl_grad)

      kl_grad_p = (kl_grad * p).sum() 
      kl_hessian = torch.autograd.grad(kl_grad_p, self.policy.parameters())
      kl_hessian = self.flat_grad(kl_hessian, hessian=True)
      return kl_hessian + p * damping_coeff