From d1dd5d578e0cbdee651eeb125dcb78e36b9cf322 Mon Sep 17 00:00:00 2001 From: Kingsley Kuan Date: Fri, 23 Oct 2020 03:39:34 +0800 Subject: [PATCH] [RLlib] Fix PyTorch A3C / A2C loss function using mixed reduced sum / mean (#11449) --- rllib/agents/a3c/a3c_torch_policy.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/rllib/agents/a3c/a3c_torch_policy.py b/rllib/agents/a3c/a3c_torch_policy.py index c91874b03..aef230c4b 100644 --- a/rllib/agents/a3c/a3c_torch_policy.py +++ b/rllib/agents/a3c/a3c_torch_policy.py @@ -18,11 +18,13 @@ def actor_critic_loss(policy, model, dist_class, train_batch): values = model.value_function() dist = dist_class(logits, model) log_probs = dist.logp(train_batch[SampleBatch.ACTIONS]) - policy.entropy = dist.entropy().mean() + policy.entropy = dist.entropy().sum() policy.pi_err = -train_batch[Postprocessing.ADVANTAGES].dot( log_probs.reshape(-1)) - policy.value_err = nn.functional.mse_loss( - values.reshape(-1), train_batch[Postprocessing.VALUE_TARGETS]) + policy.value_err = torch.sum( + torch.pow( + values.reshape(-1) - train_batch[Postprocessing.VALUE_TARGETS], + 2.0)) overall_err = sum([ policy.pi_err, policy.config["vf_loss_coeff"] * policy.value_err,