fix imageio

corl-team · Mar 5, 2024 · 9b5d6df · 9b5d6df
1 parent 27c5c72
commit 9b5d6df
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/xminigrid/experimental/img_obs.py b/src/xminigrid/experimental/img_obs.py
@@ -54,7 +54,7 @@ def build_cache(tiles: np.ndarray, tile_size: int = 32) -> tuple[np.ndarray, np.
     TILE_CACHE = jnp.asarray(TILE_CACHE).reshape(-1, TILE_SIZE, TILE_SIZE, 3)
     TILE_W_AGENT_CACHE = jnp.asarray(TILE_W_AGENT_CACHE).reshape(-1, TILE_SIZE, TILE_SIZE, 3)
 
-    print("Done. Cache will be reused on consequent runs.")
+    print(f"Done. Cache is saved to {cache_path} and will be reused on consequent runs.")
     save_bz2_pickle({"tile_cache": TILE_CACHE, "tile_agent_cache": TILE_W_AGENT_CACHE}, cache_path)
 
 TILE_CACHE = load_bz2_pickle(cache_path)["tile_cache"]

diff --git a/src/xminigrid/manual_control.py b/src/xminigrid/manual_control.py
@@ -3,7 +3,7 @@
 import argparse
 import os
 
-import imageio
+import imageio.v3 as iio
 import jax
 import numpy as np
 import pygame
@@ -167,7 +167,7 @@ def close(self) -> None:
         if self.save_video:
             assert self.video_path is not None
             save_path = os.path.join(self.video_path, "manual_control_rollout.mp4")
-            imageio.mimsave(save_path, self.frames, fps=8, format="mp4")
+            iio.imwrite(save_path, self.frames, format_hint=".mp4", fps=8)
 
 
 if __name__ == "__main__":