From e7a5d8a28481a7da9dbece78ef3ccf4a6928ed3b Mon Sep 17 00:00:00 2001
From: Hiroyuki Ota <hota@flywheel.jp>
Date: Thu, 26 Dec 2024 15:20:22 +0900
Subject: [PATCH] Fix rendering of the task by selecting the correct task_id

---
 evojax/task/slimevolley.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/evojax/task/slimevolley.py b/evojax/task/slimevolley.py
index 3e12c988..c0965258 100644
--- a/evojax/task/slimevolley.py
+++ b/evojax/task/slimevolley.py
@@ -897,7 +897,8 @@ def step(self,
     @staticmethod
     def render(state: State, task_id: int = 0) -> Image:
         """Render a specified task."""
-        game = Game(state.game_state)
+        task_game_state = jax.tree.map(lambda x: x[task_id], state.game_state)
+        game = Game(task_game_state)
         canvas = game.display()
         img = Image.fromarray(canvas)
         return img