diff --git a/src/xminigrid/experimental/img_obs.py b/src/xminigrid/experimental/img_obs.py index a90f315..759bf0a 100644 --- 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 index 5ea4068..d4174f4 100644 --- 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__":