Fixed logging info_keywords in the VecMonitor class. (#730)

* Writing the additional info_keywords into the episode infos that are passed to the resulst writer. Directly taken from the non-vec version of monitor.

* Added test for monitoring info_keywords.

* Removed unnecessary step of registering the env. Not using make_vec_env, because it applies a monitor wrapper to the env.

* Reformat

Co-authored-by: Antonin Raffin <antonin.raffin@ensta.org>

docs/misc/changelog.rst

@@ -48,6 +48,7 @@ Bug Fixes:
 - Fixed a bug where the observation would be incorrectly detected as non-vectorized instead of throwing an error
 - The env checker now properly checks and warns about potential issues for continuous action spaces when the boundaries are too small or when the dtype is not float32
 - Fixed a bug in ``VecFrameStack`` with channel first image envs, where the terminal observation would be wrongly created.
+- Fixed a bug in ``VecMonitor``. The monitor did not consider the ``info_keywords`` during stepping (@ScheiklP) 
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -879,4 +880,4 @@ And all the contributors:
 @ShangqunYu @PierreExeter @JacopoPan @ltbd78 @tom-doerr @Atlis @liusida @09tangriro @amy12xx @juancroldan
 @benblack769 @bstee615 @c-rizz @skandermoalla @MihaiAnca13 @davidblom603 @ayeright @cyprienc
 @wkirgsn @AechPro @CUN-bjy @batu @IljaAvadiev @timokau @kachayev @cleversonahum
-@eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove
+@eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove @ScheiklP

stable_baselines3/common/vec_env/vec_monitor.py

@@ -83,6 +83,8 @@ class VecMonitor(VecEnvWrapper):
                 episode_return = self.episode_returns[i]
                 episode_length = self.episode_lengths[i]
                 episode_info = {"r": episode_return, "l": episode_length, "t": round(time.time() - self.t_start, 6)}
+                for key in self.info_keywords:
+                    episode_info[key] = info[key]
                 info["episode"] = episode_info
                 self.episode_count += 1
                 self.episode_returns[i] = 0

tests/test_vec_monitor.py

@@ -1,3 +1,4 @@
+import csv
 import json
 import os
 import uuid
@@ -7,6 +8,7 @@ import pandas
 import pytest
 
 from stable_baselines3 import PPO
+from stable_baselines3.common.envs.bit_flipping_env import BitFlippingEnv
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.monitor import Monitor, get_monitor_files, load_results
 from stable_baselines3.common.vec_env import DummyVecEnv, VecMonitor, VecNormalize
@@ -45,6 +47,37 @@ def test_vec_monitor(tmp_path):
     os.remove(monitor_file)
 
 
+def test_vec_monitor_info_keywords(tmp_path):
+    """
+    Test loggig `info_keywords` in the `VecMonitor` wrapper
+    """
+    monitor_file = os.path.join(str(tmp_path), f"stable_baselines-test-{uuid.uuid4()}.monitor.csv")
+
+    env = DummyVecEnv([lambda: BitFlippingEnv()])
+
+    monitor_env = VecMonitor(env, info_keywords=("is_success",), filename=monitor_file)
+
+    monitor_env.reset()
+    total_steps = 1000
+    for _ in range(total_steps):
+        _, _, dones, infos = monitor_env.step([monitor_env.action_space.sample()])
+        if dones[0]:
+            assert "is_success" in infos[0]["episode"]
+
+    monitor_env.close()
+
+    with open(monitor_file, "rt") as f:
+        reader = csv.reader(f)
+        for i, line in enumerate(reader):
+            if i == 0 or i == 1:
+                continue
+            else:
+                assert len(line) == 4, "Incorrect keys in monitor logline"
+                assert line[3] in ["False", "True"], "Incorrect value in monitor logline"
+
+    os.remove(monitor_file)
+
+
 def test_vec_monitor_load_results(tmp_path):
     """
     test load_results on log files produced by the monitor wrapper