dont run cv2.imshow in headless environments

neuralplayground/arenas/discritized_objects.py

@@ -406,7 +406,7 @@ def plot_trajectory(
         else:
             return ax
 
-    def render(self, history_length=30):
+    def render(self, history_length=30, display=True):
         """Render the environment live through iterations"""
         f, ax = plt.subplots(1, 1, figsize=(8, 6))
         canvas = FigureCanvas(f)
@@ -416,5 +416,6 @@ def render(self, history_length=30):
         image = np.frombuffer(canvas.tostring_rgb(), dtype="uint8")
         image = image.reshape(f.canvas.get_width_height()[::-1] + (3,))
         print(image.shape)
-        cv2.imshow("2D_env", image)
-        cv2.waitKey(10)
+        if display:
+            cv2.imshow("2D_env", image)
+            cv2.waitKey(10)

neuralplayground/arenas/simple2d.py

@@ -347,7 +347,7 @@ def plot_trajectory(
         else:
             return ax
 
-    def render(self, history_length=30):
+    def render(self, history_length=30, display=True):
         """Render the environment live through iterations as in OpenAI gym"""
         f, ax = plt.subplots(1, 1, figsize=(8, 6))
         canvas = FigureCanvas(f)
@@ -357,5 +357,6 @@ def render(self, history_length=30):
         image = np.frombuffer(canvas.tostring_rgb(), dtype="uint8")
         image = image.reshape(f.canvas.get_width_height()[::-1] + (3,))
         print(image.shape)
-        cv2.imshow("2D_env", image)
-        cv2.waitKey(10)
+        if display:
+            cv2.imshow("2D_env", image)
+            cv2.waitKey(10)

tests/arena_exp_test.py

@@ -51,7 +51,7 @@ def test_agent_interaction(self, init_env):
             action = agent.act(obs)
             # Run environment for given action
             obs, state, reward = init_env[0].step(action)
-            init_env[0].render()
+            init_env[0].render(display=False)
         init_env[0].plot_trajectory()