diff --git a/cooperative_cuisine/pygame_2d_vis/drawing.py b/cooperative_cuisine/pygame_2d_vis/drawing.py
index 451129494998a0d2da363c157a571d885eed8ca9..88442ea8356f22c926ddf789f1c8d03f06c1814f 100644
--- a/cooperative_cuisine/pygame_2d_vis/drawing.py
+++ b/cooperative_cuisine/pygame_2d_vis/drawing.py
@@ -937,8 +937,10 @@ class Visualizer:
         pygame.image.save(screen, filename)
 
     def get_state_image(self, state: dict,
-                        cache_flags: CacheFlags = CacheFlags.COUNTERS | CacheFlags.BACKGROUND) -> npt.NDArray[np.uint8]:
-        screen = self.draw_gamescreen(state, [0 for _ in state["players"]], cache_flags=cache_flags)
+                        cache_flags: CacheFlags = CacheFlags.COUNTERS | CacheFlags.BACKGROUND,
+                        controlled_players: list[int] = None,
+                        ) -> npt.NDArray[np.uint8]:
+        screen = self.draw_gamescreen(state, [0] if controlled_players is None else controlled_players, cache_flags=cache_flags)
         return pygame.surfarray.pixels3d(screen)
 
     def draw_recipe_image(