feat(overcooked): Fix viz tile rotation

jaxmarl/viz/grid_rendering_v2.py

@@ -32,7 +32,7 @@ def _mask_fn(y, x):
         return fn(xf, yf)
 
     ys, xs = jnp.indices(img.shape[:2])
-    mask = jax.vmap(jax.vmap(_mask_fn, in_axes=0), in_axes=1)(ys, xs)
+    mask = jax.vmap(jax.vmap(_mask_fn))(ys, xs)
 
     color_img = jnp.full_like(img, color)
     return jnp.where(mask[:, :, None], color_img, img)

jaxmarl/viz/overcooked_v2_visualizer.py

@@ -78,11 +78,17 @@ def animate(self, state_seq, filename="animation.gif", agent_view_size=None):
             state_seq, agent_view_size
         )
         # print("frame_seq", frame_seq)
-        print("frame_seq.shape", frame_seq.shape)
-        print("frame_seq.dtype", frame_seq.dtype)
+        # print("frame_seq.shape", frame_seq.shape)
+        # print("frame_seq.dtype", frame_seq.dtype)
 
         imageio.mimsave(filename, frame_seq, "GIF", duration=0.5)
 
+    def render_sequence(self, state_seq, agent_view_size=None):
+        frame_seq = jax.vmap(self._render_state, in_axes=(0, None))(
+            state_seq, agent_view_size
+        )
+        return frame_seq
+
     @partial(jax.jit, static_argnums=(0, 2))
     def _render_state(self, state, agent_view_size=None):
         """
@@ -414,15 +420,15 @@ def _render_grid(
     ):
         img_grid = jax.vmap(jax.vmap(self._render_tile))(grid, highlight_mask)
 
-        print("img_grid", img_grid.shape)
+        # print("img_grid", img_grid.shape)
 
         grid_rows, grid_cols, tile_height, tile_width, channels = img_grid.shape
 
         big_image = img_grid.transpose(0, 2, 1, 3, 4).reshape(
             grid_rows * tile_height, grid_cols * tile_width, channels
         )
 
-        print("big_image", big_image.shape)
+        # print("big_image", big_image.shape)
 
         return big_image