Fix vtrace (#286)

* Checking for vtrace before adding advantage stats to summary * Fixed and formatted --------- Co-authored-by: alex404 <[email protected]>
alex-petrenko · Oct 14, 2023 · 37db90f · 37db90f
1 parent 7e1e695
commit 37db90f
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/sample_factory/algo/learning/learner.py b/sample_factory/algo/learning/learner.py
@@ -866,10 +866,12 @@ def _record_summaries(self, train_loop_vars) -> AttrDict:
         stats.act_min = var.mb.actions.min()
         stats.act_max = var.mb.actions.max()
 
-        stats.adv_min = var.mb.advantages.min()
-        stats.adv_max = var.mb.advantages.max()
-        stats.adv_std = var.adv_std
-        stats.adv_mean = var.adv_mean
+        if "adv_mean" in stats:
+            stats.adv_min = var.mb.advantages.min()
+            stats.adv_max = var.mb.advantages.max()
+            stats.adv_std = var.adv_std
+            stats.adv_mean = var.adv_mean
+
         stats.max_abs_logprob = torch.abs(var.mb.action_logits).max()
 
         if hasattr(var.action_distribution, "summaries"):