Hi Thomas, thanks very much for your reply! I’m already sampling mini-batches (of size 32) from my replay buffer but this is still incredibly slow (even though I’m only using a toy environment at the moment). Perhaps to make things more concrete, my (pseudo-ish) code for computing the scores is as follows:
scores = []

log_probs = self.policy(states).log_prob(actions)

for params in self.actor.parameters():
    scores.append(torch.stack([torch.autograd.grad(outputs=l_p, inputs=params, retain_graph=True)[0] for l_p in log_probs]))

It may well be that this is the only/fastest way of doing things (although I doubt that!), but in which case I suppose I just need to perform this computation as few times as I possible can during learning.
Any further suggestions or comments are greatly appreciated