From 4b5e5d5806c89303a625386a4a6aae200327fdb5 Mon Sep 17 00:00:00 2001
From: bragajj <39658109+bragajj@users.noreply.github.com>
Date: Mon, 24 Jan 2022 09:28:46 -0800
Subject: [PATCH] Support for VecVideoRecorder (#62)

* Support for VecVideoRecorder

Support for VecVideoRecorder
 https://stable-baselines3.readthedocs.io/en/master/guide/vec_envs.html?highlight=video#vecvideorecorder

* VecVideoRecorder update

* Update env.py

Frames per second fixed 24 from to 15 fps
---
 procgen/env.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/procgen/env.py b/procgen/env.py
index 1addfe6f..b7b7590a 100644
--- a/procgen/env.py
+++ b/procgen/env.py
@@ -244,10 +244,22 @@ def __init__(
                 "distribution_mode": distribution_mode,
             }
         super().__init__(num, env_name, options, **kwargs)
+        
+        
+class ToBaselinesVecEnv(gym3.ToBaselinesVecEnv):
+    metadata = {
+        'render.modes': ['human', 'rgb_array'],
+        'video.frames_per_second' : 15
+    }
+    def render(self, mode="human"):
+        info = self.env.get_info()[0]
+        _, ob, _ = self.env.observe()
+        if mode == "rgb_array":
+            if "rgb" in info:
+                return info["rgb"]
+            else:
+                return ob['rgb'][0]        
 
 
 def ProcgenEnv(num_envs, env_name, **kwargs):
-    """
-    Baselines VecEnv interface for Procgen
-    """
-    return gym3.ToBaselinesVecEnv(ProcgenGym3Env(num=num_envs, env_name=env_name, **kwargs))
\ No newline at end of file
+    return ToBaselinesVecEnv(ProcgenGym3Env(num=num_envs, env_name=env_name, **kwargs))