Merge branch 'ts/add_fruit_tree_env_render' of https://github.com/tom…

…ekster/MO-Gymnasium into ts/add_fruit_tree_env_render
Farama-Foundation · Jan 27, 2024 · 36f2ce6 · 36f2ce6
2 parents 70b3fe8 + 151f8f4
commit 36f2ce6
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/mo_gymnasium/envs/fruit_tree/fruit_tree.py b/mo_gymnasium/envs/fruit_tree/fruit_tree.py
@@ -452,7 +452,14 @@ def render(self):
             pygame.display.update()
             self.clock.tick(self.metadata["render_fps"])
         elif self.render_mode == "rgb_array":
-            return np.transpose(np.array(pygame.surfarray.pixels3d(self.window)), axes=(1, 0, 2))
+            return np.transpose(np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2))
+
+        background = pygame.Surface(self.window_size)
+        background.fill((255, 255, 255))  # White
+
+        background.blit(canvas, (0, self.top_margin))
+
+        self.window.blit(background, (0, 0))
 
 
 if __name__ == "__main__":