diff --git a/README.md b/README.md
index 8f7b836..093fec3 100644
--- a/README.md
+++ b/README.md
@@ -64,7 +64,7 @@ Like with other `gymnasium` environments, it's very easy to use `flappy-bird-gym
 Simply import the package and create the environment with the `make` function.
 Take a look at the sample code below:
 
-```
+```python
 import time
 import flappy_bird_gymnasium
 import gymnasium
diff --git a/flappy_bird_gymnasium/envs/flappy_bird_env_rgb.py b/flappy_bird_gymnasium/envs/flappy_bird_env_rgb.py
index d8b74a7..06bc34b 100644
--- a/flappy_bird_gymnasium/envs/flappy_bird_env_rgb.py
+++ b/flappy_bird_gymnasium/envs/flappy_bird_env_rgb.py
@@ -66,7 +66,7 @@ def __init__(
         pipe_gap: int = 100,
         bird_color: str = "yellow",
         pipe_color: str = "green",
-        render_mode=None,
+        render_mode: Optional[str] = None,
         background: Optional[str] = None,
     ) -> None:
         self.action_space = gymnasium.spaces.Discrete(2)
diff --git a/flappy_bird_gymnasium/envs/flappy_bird_env_simple.py b/flappy_bird_gymnasium/envs/flappy_bird_env_simple.py
index 3435017..d8634f2 100644
--- a/flappy_bird_gymnasium/envs/flappy_bird_env_simple.py
+++ b/flappy_bird_gymnasium/envs/flappy_bird_env_simple.py
@@ -25,6 +25,7 @@
 numerical information about the game's state as observations.
 """
 
+import time
 from typing import Dict, Optional, Tuple, Union
 
 import gymnasium
@@ -70,16 +71,17 @@ class FlappyBirdEnvSimple(gymnasium.Env):
             be drawn.
     """
 
-    metadata = {"render_modes": ["human"], "render_fps": 30}
+    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 30}
 
     def __init__(
         self,
         screen_size: Tuple[int, int] = (288, 512),
-        audio_on: bool = True,
+        audio_on: bool = False,
         normalize_obs: bool = True,
         pipe_gap: int = 100,
         bird_color: str = "yellow",
         pipe_color: str = "green",
+        render_mode: Optional[str] = None,
         background: Optional[str] = "day",
     ) -> None:
         self.action_space = gymnasium.spaces.Discrete(2)
@@ -98,6 +100,18 @@ def __init__(
         self._pipe_color = pipe_color
         self._bg_type = background
 
+        assert render_mode is None or render_mode in self.metadata["render_modes"]
+        self.render_mode = render_mode
+
+        if render_mode is not None:
+            self._renderer = FlappyBirdRenderer(
+                screen_size=self._screen_size,
+                audio_on=audio_on,
+                bird_color=bird_color,
+                pipe_color=pipe_color,
+                background=background,
+            )
+
     def _get_observation(self):
         pipes = []
         for up_pipe, low_pipe in zip(self._game.upper_pipes, self._game.lower_pipes):
@@ -193,19 +207,16 @@ def set_color(self, color):
 
     def render(self) -> None:
         """Renders the next frame."""
-        if self._renderer is None:
-            self._renderer = FlappyBirdRenderer(
-                screen_size=self._screen_size,
-                audio_on=self._audio_on,
-                bird_color=self._bird_color,
-                pipe_color=self._pipe_color,
-                background=self._bg_type,
-            )
-            self._renderer.game = self._game
-            self._renderer.make_display()
-
         self._renderer.draw_surface(show_score=True)
-        self._renderer.update_display()
+        if self.render_mode == "rgb_array":
+            # Flip the image to retrieve a correct aspect
+            return np.transpose(pygame.surfarray.array3d(self._renderer.surface), axes=(1, 0, 2))
+        else:
+            if self._renderer.display is None:
+                self._renderer.make_display()
+
+            self._renderer.update_display()
+            time.sleep(1 / self.metadata["render_fps"])
 
     def close(self):
         """Closes the environment."""
diff --git a/flappy_bird_gymnasium/envs/renderer.py b/flappy_bird_gymnasium/envs/renderer.py
index 1781e4e..74b83c4 100644
--- a/flappy_bird_gymnasium/envs/renderer.py
+++ b/flappy_bird_gymnasium/envs/renderer.py
@@ -182,6 +182,7 @@ def update_display(self) -> None:
                 "call the `make_display()` method."
             )
 
+        pygame.event.get()
         self.display.blit(self.surface, [0, 0])
         pygame.display.update()
 
diff --git a/setup.py b/setup.py
index e7d5cff..69a0a00 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@
 
 import setuptools
 
-_VERSION = "0.2.1"
+_VERSION = "0.2.2"
 
 # Short description.
 short_description = "A Gymnasium environment for the Flappy Bird game."