From f86c6c736f788650ef5723f59b2729e290dc981e Mon Sep 17 00:00:00 2001
From: Dominik Roth <dominik.roth.dev@gmail.com>
Date: Thu, 26 Jan 2023 18:15:06 +0100
Subject: [PATCH] Hack is back

---
 metastable_baselines/misc/rollout_buffer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/metastable_baselines/misc/rollout_buffer.py b/metastable_baselines/misc/rollout_buffer.py
index 846d60a..1fea969 100644
--- a/metastable_baselines/misc/rollout_buffer.py
+++ b/metastable_baselines/misc/rollout_buffer.py
@@ -236,6 +236,8 @@ class GaussianRolloutCollectorAuxclass():
             if callback.on_step() is False:
                 return False
 
+            if isinstance(env, DummyVecEnv):
+                self._update_info_buffer(infos[0])
             self._update_info_buffer(infos)
             n_steps += 1