I’m relatively new to Pytorch and am trying to implement a natural actor-critic (NAC) algorithm where my policy pi is a DNN. See algorithm 2 in Thomas (2014) 1 for a simple example. In order to do this, I want to compute the natural gradient x by minimising the following loss, where nu is the on-policy state-action distribution, and A is the advantage function (which I already have):
Screenshot 2020-09-28 at 19.31.58733×131 12.1 KB
My question is: How do I best (most efficiently) compute the score function (the derivative of the log of the policy w.r.t. its parameters theta)?
I have a naive implementation already, but it involves computing this score function for each individual state-action sample using a single application of torch.autograd.grad, and is therefore horrendously slow. I am aware that NAC algorithms have been used before for fairly high-dimensional control tasks, but have been unable to find an existing implementation online.
Any help would be hugely appreciated, thank you!