fix tqdm issue (thu-ml#481)

BFAnas · Nov 18, 2021 · 38192a9 · 38192a9
1 parent d8b5b30
commit 38192a9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tianshou/trainer/offpolicy.py b/tianshou/trainer/offpolicy.py
@@ -117,8 +117,8 @@ def offpolicy_trainer(
                 env_step += int(result["n/st"])
                 t.update(result["n/st"])
                 logger.log_train_data(result, env_step)
-                last_rew = result['rew'] if 'rew' in result else last_rew
-                last_len = result['len'] if 'len' in result else last_len
+                last_rew = result['rew'] if result["n/ep"] > 0 else last_rew
+                last_len = result['len'] if result["n/ep"] > 0 else last_len
                 data = {
                     "env_step": str(env_step),
                     "rew": f"{last_rew:.2f}",

diff --git a/tianshou/trainer/onpolicy.py b/tianshou/trainer/onpolicy.py
@@ -125,8 +125,8 @@ def onpolicy_trainer(
                 env_step += int(result["n/st"])
                 t.update(result["n/st"])
                 logger.log_train_data(result, env_step)
-                last_rew = result['rew'] if 'rew' in result else last_rew
-                last_len = result['len'] if 'len' in result else last_len
+                last_rew = result['rew'] if result["n/ep"] > 0 else last_rew
+                last_len = result['len'] if result["n/ep"] > 0 else last_len
                 data = {
                     "env_step": str(env_step),
                     "rew": f"{last_rew:.2f}",