fix render popping when render is called internally for Wrappers (#2998)

* fix render popping when render is called internally

* move _render workaround from Wrapper to PixelObservation

Author: younik

gym/core.py

@@ -426,7 +426,9 @@ def reset(self, **kwargs) -> Union[ObsType, Tuple[ObsType, dict]]:
         """Resets the environment with kwargs."""
         return self.env.reset(**kwargs)
 
-    def render(self, *args, **kwargs):
+    def render(
+        self, *args, **kwargs
+    ) -> Optional[Union[RenderFrame, List[RenderFrame]]]:
         """Renders the environment."""
         return self.env.render(*args, **kwargs)
 

gym/wrappers/pixel_observation.py

@@ -81,6 +81,7 @@ def __init__(
 
         # Avoid side-effects that occur when render_kwargs is manipulated
         render_kwargs = copy.deepcopy(render_kwargs)
+        self.render_history = []
 
         if render_kwargs is None:
             render_kwargs = {}
@@ -135,7 +136,7 @@ def __init__(
         self.env.reset()
         pixels_spaces = {}
         for pixel_key in pixel_keys:
-            pixels = self.env.render(**render_kwargs[pixel_key])
+            pixels = self._render(**render_kwargs[pixel_key])
             pixels: np.ndarray = pixels[-1] if isinstance(pixels, List) else pixels
 
             if not hasattr(pixels, "dtype") or not hasattr(pixels, "shape"):
@@ -184,10 +185,24 @@ def _add_pixel_observation(self, wrapped_observation):
             observation[STATE_KEY] = wrapped_observation
 
         pixel_observations = {
-            pixel_key: self.env.render(**self._render_kwargs[pixel_key])
+            pixel_key: self._render(**self._render_kwargs[pixel_key])
             for pixel_key in self._pixel_keys
         }
 
         observation.update(pixel_observations)
 
         return observation
+
+    def render(self, *args, **kwargs):
+        """Renders the environment."""
+        render = self.env.render(*args, **kwargs)
+        if isinstance(render, list):
+            render = self.render_history + render
+            self.render_history = []
+        return render
+
+    def _render(self, *args, **kwargs):
+        render = self.env.render(*args, **kwargs)
+        if isinstance(render, list):
+            self.render_history += render
+        return render