I’m trying to implement the trust region update from Sample Efficient Actor-Critic with Experience Replay 33 (Algorithm 2, Appendix A). I have Chainer code as reference 16, which I’ve translated as follows:
def trust_region_loss(model, ref_model, distribution, ref_distribution, loss, threshold):
  # Compute gradients from original loss
  loss.backward(retain_variables=True)
  g = [param.grad.clone() for param in model.parameters()]
  model.zero_grad()

  # KL divergence k ← ∇θ0∙DKL[π(∙|s_i; θ_a) || π(∙|s_i; θ)]
  kl = distribution * (distribution.log() - ref_distribution.log())
  # Compute gradients from (negative) KL loss (increases KL divergence)
  (-kl.mean(1)).backward(retain_variables=True)
  k = [param.grad.clone() for param in model.parameters()]
  model.zero_grad()

  # Compute dot products of gradients
  k_dot_g = sum(torch.sum(k_p * g_p) for k_p, g_p in zip(k, g))
  k_dot_k = sum(torch.sum(k_p ** 2) for k_p in k)
  # Compute trust region update
  trust_factor = k_dot_k.data[0] > 0 and (k_dot_g - threshold) / k_dot_k or Variable(torch.zeros(1))
  trust_update = [g_p - trust_factor.expand_as(k_p) * k_p for g_p, k_p in zip(g, k)]
  trust_loss = 0
  for param, trust_update_p in zip(model.parameters(), trust_update):
    trust_loss += (param * trust_update_p).sum()

However, the eventual .backward() call on this loss results in RuntimeError: calling backward on a volatile variable, which I assume is due to trying to use gradients in this way. Adding trust_loss.volatile = False allows the backwards pass to happen, but the network doesn’t seem to improve (could be a matter of hyperparameter tuning here). Is this the correct approach?