Print episode stats iff ep ended during collect

dantp-ai · Jan 3, 2024 · a6a13ce · a6a13ce
1 parent 1468cdf
commit a6a13ce
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/offline/test_bcq.py b/test/offline/test_bcq.py
@@ -207,7 +207,8 @@ def watch():
         policy.eval()
         collector = Collector(policy, env)
         res = collector.collect(n_episode=1, render=args.render)
-        print(f"Final reward: {res.returns_stat.mean}, length: {res.lens_stat.mean}")
+        if res.returns_stat and res.lens_stat:
+            print(f"Final reward: {res.returns_stat.mean}, length: {res.lens_stat.mean}")
 
 
 if __name__ == "__main__":