Prevent lineplot crash if no complete train trajectories

Kaixhin · Aug 6, 2021 · aff0689 · aff0689
1 parent cc1c621
commit aff0689
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/main.py b/main.py
@@ -145,7 +145,7 @@ def main(cfg: DictConfig) -> None:
       if cfg.algorithm == 'BC':
         lineplot(range(cfg.evaluation.interval, cfg.steps + 1, cfg.evaluation.interval), metrics['test_returns'] * (cfg.steps // cfg.evaluation.interval), 'test_returns')
         break
-      else:
+      elif len(metrics['train_returns']) > 0:  # Plot train returns if any
         lineplot(metrics['train_steps'], metrics['train_returns'], 'train_returns')
     elif cfg.algorithm == 'BC' and cfg.check_time_usage: break