Repository: openai/gym

Issue:
[Bug Report] FrozenLake RGB rendering returns None
FrozenLake `render()` method in mode human always returns `None` and lacks an `rgb_array` render mode. For reference, Blackjack has recently been converted to pygame and implements these features correctly.

@younik I believe you converted FrozenLake to pygame? 



Discussion:

Resolving patch:
diff --git a/gym/envs/toy_text/taxi.py b/gym/envs/toy_text/taxi.py
--- a/gym/envs/toy_text/taxi.py
+++ b/gym/envs/toy_text/taxi.py
@@ -1,6 +1,6 @@
-import sys
 from contextlib import closing
 from io import StringIO
+from os import path
 from typing import Optional
 
 import numpy as np
@@ -16,6 +16,7 @@
     "|Y| : |B: |",
     "+---------+",
 ]
+WINDOW_SIZE = (550, 350)
 
 
 class TaxiEnv(Env):
@@ -82,12 +83,6 @@ class TaxiEnv(Env):
     - +20 delivering passenger.
     - -10  executing "pickup" and "drop-off" actions illegally.
 
-    ### Rendering
-    - blue: passenger
-    - magenta: destination
-    - yellow: empty taxi
-    - green: full taxi
-    - other letters (R, G, Y and B): locations for passengers and destinations
     state space is represented by:
     (taxi_row, taxi_col, passenger_location, destination)
 
@@ -104,12 +99,13 @@ class TaxiEnv(Env):
     * v0: Initial versions release
     """
 
-    metadata = {"render_modes": ["human", "ansi"], "render_fps": 4}
+    metadata = {"render_modes": ["human", "ansi", "rgb_array"], "render_fps": 4}
 
     def __init__(self):
         self.desc = np.asarray(MAP, dtype="c")
 
         self.locs = locs = [(0, 0), (0, 4), (4, 0), (4, 3)]
+        self.locs_colors = [(255, 0, 0), (0, 255, 0), (255, 255, 0), (0, 0, 255)]
 
         num_states = 500
         num_rows = 5
@@ -168,6 +164,21 @@ def __init__(self):
         self.action_space = spaces.Discrete(num_actions)
         self.observation_space = spaces.Discrete(num_states)
 
+        # pygame utils
+        self.window = None
+        self.clock = None
+        self.cell_size = (
+            WINDOW_SIZE[0] / self.desc.shape[1],
+            WINDOW_SIZE[1] / self.desc.shape[0],
+        )
+        self.taxi_imgs = None
+        self.taxi_orientation = 0
+        self.passenger_img = None
+        self.destination_img = None
+        self.median_horiz = None
+        self.median_vert = None
+        self.background_img = None
+
     def encode(self, taxi_row, taxi_col, pass_loc, dest_idx):
         # (5) 5, 5, 4
         i = taxi_row
@@ -209,16 +220,150 @@ def reset(
         super().reset(seed=seed)
         self.s = categorical_sample(self.initial_state_distrib, self.np_random)
         self.lastaction = None
+        self.taxi_orientation = 0
         if not return_info:
             return int(self.s)
         else:
             return int(self.s), {"prob": 1}
 
     def render(self, mode="human"):
-        outfile = StringIO() if mode == "ansi" else sys.stdout
+        if mode == "ansi":
+            return self._render_text()
+        else:
+            return self._render_gui(mode)
+
+    def _render_gui(self, mode):
+        import pygame  # dependency to pygame only if rendering with human
+
+        if self.window is None:
+            pygame.init()
+            pygame.display.set_caption("Taxi")
+            if mode == "human":
+                self.window = pygame.display.set_mode(WINDOW_SIZE)
+            else:  # rgb_array
+                self.window = pygame.Surface(WINDOW_SIZE)
+        if self.clock is None:
+            self.clock = pygame.time.Clock()
+        if self.taxi_imgs is None:
+            file_names = [
+                path.join(path.dirname(__file__), "img/cab_front.png"),
+                path.join(path.dirname(__file__), "img/cab_rear.png"),
+                path.join(path.dirname(__file__), "img/cab_right.png"),
+                path.join(path.dirname(__file__), "img/cab_left.png"),
+            ]
+            self.taxi_imgs = [
+                pygame.transform.scale(pygame.image.load(file_name), self.cell_size)
+                for file_name in file_names
+            ]
+        if self.passenger_img is None:
+            file_name = path.join(path.dirname(__file__), "img/passenger.png")
+            self.passenger_img = pygame.transform.scale(
+                pygame.image.load(file_name), self.cell_size
+            )
+        if self.destination_img is None:
+            file_name = path.join(path.dirname(__file__), "img/hotel.png")
+            self.destination_img = pygame.transform.scale(
+                pygame.image.load(file_name), self.cell_size
+            )
+            self.destination_img = self.destination_img.convert_alpha()
+            self.destination_img.set_alpha(170)
+        if self.median_horiz is None:
+            file_names = [
+                path.join(path.dirname(__file__), "img/gridworld_median_left.png"),
+                path.join(path.dirname(__file__), "img/gridworld_median_horiz.png"),
+                path.join(path.dirname(__file__), "img/gridworld_median_right.png"),
+            ]
+            self.median_horiz = [
+                pygame.transform.scale(pygame.image.load(file_name), self.cell_size)
+                for file_name in file_names
+            ]
+        if self.median_vert is None:
+            file_names = [
+                path.join(path.dirname(__file__), "img/gridworld_median_top.png"),
+                path.join(path.dirname(__file__), "img/gridworld_median_vert.png"),
+                path.join(path.dirname(__file__), "img/gridworld_median_bottom.png"),
+            ]
+            self.median_vert = [
+                pygame.transform.scale(pygame.image.load(file_name), self.cell_size)
+                for file_name in file_names
+            ]
+        if self.background_img is None:
+            file_name = path.join(path.dirname(__file__), "img/taxi_background.png")
+            self.background_img = pygame.transform.scale(
+                pygame.image.load(file_name), self.cell_size
+            )
+
+        desc = self.desc
+
+        for y in range(0, desc.shape[0]):
+            for x in range(0, desc.shape[1]):
+                cell = (x * self.cell_size[0], y * self.cell_size[1])
+                self.window.blit(self.background_img, cell)
+                if desc[y][x] == b"|" and (y == 0 or desc[y - 1][x] != b"|"):
+                    self.window.blit(self.median_vert[0], cell)
+                elif desc[y][x] == b"|" and (
+                    y == desc.shape[0] - 1 or desc[y + 1][x] != b"|"
+                ):
+                    self.window.blit(self.median_vert[2], cell)
+                elif desc[y][x] == b"|":
+                    self.window.blit(self.median_vert[1], cell)
+                elif desc[y][x] == b"-" and (x == 0 or desc[y][x - 1] != b"-"):
+                    self.window.blit(self.median_horiz[0], cell)
+                elif desc[y][x] == b"-" and (
+                    x == desc.shape[1] - 1 or desc[y][x + 1] != b"-"
+                ):
+                    self.window.blit(self.median_horiz[2], cell)
+                elif desc[y][x] == b"-":
+                    self.window.blit(self.median_horiz[1], cell)
+
+        for cell, color in zip(self.locs, self.locs_colors):
+            color_cell = pygame.Surface(self.cell_size)
+            color_cell.set_alpha(128)
+            color_cell.fill(color)
+            loc = self.get_surf_loc(cell)
+            self.window.blit(color_cell, (loc[0], loc[1] + 10))
+
+        taxi_row, taxi_col, pass_idx, dest_idx = self.decode(self.s)
+
+        if pass_idx < 4:
+            self.window.blit(self.passenger_img, self.get_surf_loc(self.locs[pass_idx]))
 
-        out = self.desc.copy().tolist()
-        out = [[c.decode("utf-8") for c in line] for line in out]
+        if self.lastaction in [0, 1, 2, 3]:
+            self.taxi_orientation = self.lastaction
+        dest_loc = self.get_surf_loc(self.locs[dest_idx])
+        taxi_location = self.get_surf_loc((taxi_row, taxi_col))
+
+        if dest_loc[1] <= taxi_location[1]:
+            self.window.blit(
+                self.destination_img,
+                (dest_loc[0], dest_loc[1] - self.cell_size[1] // 2),
+            )
+            self.window.blit(self.taxi_imgs[self.taxi_orientation], taxi_location)
+        else:  # change blit order for overlapping appearance
+            self.window.blit(self.taxi_imgs[self.taxi_orientation], taxi_location)
+            self.window.blit(
+                self.destination_img,
+                (dest_loc[0], dest_loc[1] - self.cell_size[1] // 2),
+            )
+
+        if mode == "human":
+            pygame.display.update()
+            self.clock.tick(self.metadata["render_fps"])
+        else:  # rgb_array
+            return np.transpose(
+                np.array(pygame.surfarray.pixels3d(self.window)), axes=(1, 0, 2)
+            )
+
+    def get_surf_loc(self, map_loc):
+        return (map_loc[1] * 2 + 1) * self.cell_size[0], (
+            map_loc[0] + 1
+        ) * self.cell_size[1]
+
+    def _render_text(self):
+        desc = self.desc.copy().tolist()
+        outfile = StringIO()
+
+        out = [[c.decode("utf-8") for c in line] for line in desc]
         taxi_row, taxi_col, pass_idx, dest_idx = self.decode(self.s)
 
         def ul(x):
@@ -247,7 +392,12 @@ def ul(x):
         else:
             outfile.write("\n")
 
-        # No need to return anything for human
-        if mode != "human":
-            with closing(outfile):
-                return outfile.getvalue()
+        with closing(outfile):
+            return outfile.getvalue()
+
+    def close(self):
+        if self.window is not None:
+            import pygame
+
+            pygame.display.quit()
+            pygame.quit()


Test patch: