Gif for gymnasium env

docs/basics/userguide/visualization.md

@@ -39,6 +39,41 @@ env.disable_rendering()
 
 The tool `save_gif` works by collecting all the frames generated by the steps of `env` for each of the steps done after `env.enable_rendering()`.
 
+<span>&#9888;&#9888;&#9888;</span> **warning :** Gymnasium<span>&#9888;&#9888;&#9888;</span>
+
+Be carful with [gymnasium](https://gymnasium.farama.org/) env.
+If you want to use the `save_gif` tool with a gymnasium environment, you need to :
+ - Use the {mod}`rlberry.envs.gym_make` function to create the (wrapped) environment.
+ - Use `render_mode = "rgb_array"` as parameter to create the environment. More information [here](https://gymnasium.farama.org/api/env/#gymnasium.Env.render).
+
+An example here :
+```python
+from rlberry.envs import gym_make
+from rlberry.seeding import Seeder
+import numpy as np
+
+
+seeder = Seeder(123)
+env = gym_make("MountainCar-v0", render_mode="rgb_array")
+env.reseed(seeder)
+
+
+env.enable_rendering()
+observation, info = env.reset()
+for tt in range(100):
+    action = np.random.randint(0, 1)
+    observation, reward, terminated, truncated, info = env.step(action)
+    done = terminated or truncated
+    if done:
+        # Warning: this will never happen in the present case.
+        break
+
+saving_path = "./visu_gymnasium_gif.gif"
+env.save_gif(saving_path)
+```
+
+![](visu_gymnasium_gif.gif)
+
 ## Plotting training data and reward curves in rlberry
 
 Training metrics in reinforcement learning are typically very unstable because there is randomness in the training process (e.g. neural networks) and the environment changes during the training. Then, if a single plot is very non-smooth due to the process variablility, it may become necessary to smooth the curve in order to be able to get any information from it.

docs/changelog.rst

@@ -7,7 +7,9 @@ Changelog
 Dev version
 -----------
 
- *nothing*
+ *PR #467*
+
+ * Allow the save as Gif for gymnasium env (make_gym) : https://github.com/rlberry-py/rlberry/issues/453
 
 
 Version 0.7.3

rlberry/envs/basewrapper.py

@@ -2,10 +2,12 @@
 from rlberry.seeding import Seeder, safe_reseed
 import numpy as np
 from rlberry.envs.interface import Model
+from rlberry.rendering import RenderInterface
 from rlberry.spaces.from_gym import convert_space_from_gym
+from rlberry.rendering.utils import video_write, gif_write
 
 
-class Wrapper(Model):
+class Wrapper(Model, RenderInterface):
     """
     Wraps a given environment, similar to OpenAI gym's wrapper [1,2] (now updated to gymnasium).
     Can also be used to wrap gym environments.
@@ -46,6 +48,7 @@ def __init__(self, env, wrap_spaces=False):
         self.env = env
         self.metadata = self.env.metadata
         self.render_mode = self.env.render_mode
+        self.frames = []
 
         if wrap_spaces:
             self.observation_space = convert_space_from_gym(self.env.observation_space)
@@ -107,11 +110,14 @@ def reseed(self, seed_seq=None):
     def reset(self, seed=None, options=None):
         if self.env.render_mode == "human":
             self.render()
+        self.frames = []
         return self.env.reset(seed=seed, options=options)
 
     def step(self, action):
         if self.render_mode == "human":
             self.render()
+        elif self.render_mode == "rgb_array":
+            self.frames.append(self.render())
         return self.env.step(action)
 
     def sample(self, state, action):
@@ -147,6 +153,17 @@ def is_generative(self):
         except Exception:
             return False
 
+    def get_video(self, **kwargs):
+        return self.frames
+
+    def save_video(self, filename, framerate=25, **kwargs):
+        video_data = self.get_video(**kwargs)
+        video_write(filename, video_data, framerate=framerate)
+
+    def save_gif(self, filename, **kwargs):
+        video_data = self.get_video(**kwargs)
+        gif_write(filename, video_data)
+
     def __repr__(self):
         return str(self)
 

rlberry/rendering/render_interface.py

@@ -139,7 +139,7 @@ def render(self, loop=True, **kwargs):
             logger.info("Rendering not enabled for the environment.")
             return 1
 
-    def get_video(self, framerate=25, **kwargs):
+    def get_video(self, **kwargs):
         # background and data
         background, data = self._get_background_and_scenes()
 
@@ -158,9 +158,9 @@ def get_video(self, framerate=25, **kwargs):
         return renderer.get_video_data()
 
     def save_video(self, filename, framerate=25, **kwargs):
-        video_data = self.get_video(framerate=framerate, **kwargs)
+        video_data = self.get_video(**kwargs)
         video_write(filename, video_data, framerate=framerate)
 
-    def save_gif(self, filename, framerate=25, **kwargs):
-        video_data = self.get_video(framerate=framerate, **kwargs)
+    def save_gif(self, filename, **kwargs):
+        video_data = self.get_video(**kwargs)
         gif_write(filename, video_data)

rlberry/rendering/tests/test_rendering_interface.py

@@ -8,6 +8,7 @@
 from rlberry_scool.envs.finite import Chain
 from rlberry_scool.envs.finite import GridWorld
 from rlberry_scool.agents.dynprog import ValueIterationAgent
+from rlberry_research.agents import RSUCBVIAgent
 from rlberry_research.envs.benchmarks.grid_exploration.four_room import FourRoom
 from rlberry_research.envs.benchmarks.grid_exploration.six_room import SixRoom
 from rlberry_research.envs.benchmarks.grid_exploration.apple_gold import AppleGold
@@ -16,6 +17,8 @@
 from rlberry.rendering import RenderInterface
 from rlberry.rendering import RenderInterface2D
 from rlberry.envs import Wrapper
+from rlberry.envs import gym_make
+from rlberry.seeding import Seeder
 
 import tempfile
 
@@ -156,7 +159,7 @@
         done = terminated or truncated
         if done:
             # Warning: this will never happen in the present case because there is no terminal state.
-            # See the doc of GridWorld for more informations on the default parameters of GridWorld.
+            # See the doc of GridWorld for more information on the default parameters of GridWorld.
             break
 
     with tempfile.TemporaryDirectory() as tmpdirname:
@@ -169,7 +172,7 @@
             pass
 
 
-##### Fonctionne uniquement si on ajoute une dépendance à ffmpeg ############
+# ##### Works only if you add a dependency to ffmpeg ############
 # @pytest.mark.xfail(sys.platform == "darwin", reason="bug with Mac with pygame")
 # @pytest.mark.parametrize("rendering_tool", RENDERING_TOOL)
 # def test_gridworld_rendering_mp4(rendering_tool):
@@ -188,11 +191,11 @@
 #         done = terminated or truncated
 #         if done:
 #             # Warning: this will never happen in the present case because there is no terminal state.
-#             # See the doc of GridWorld for more informations on the default parameters of GridWorld.
+#             # See the doc of GridWorld for more information on the default parameters of GridWorld.
 #             break
 
 #     with tempfile.TemporaryDirectory() as tmpdirname:
-#         saving_path = tmpdirname + "/test_gif.mp4"
+#         saving_path = tmpdirname + "/test_video.mp4"
 #         env.save_video(saving_path)
 #         assert os.path.isfile(saving_path)
 #         try:
@@ -219,7 +222,89 @@
         done = terminated or truncated
         if done:
             # Warning: this will never happen in the present case because there is no terminal state.
-            # See the doc of GridWorld for more informations on the default parameters of GridWorld.
+            # See the doc of GridWorld for more information on the default parameters of GridWorld.
             break
 
     env.render(loop=False)
+
+
+@pytest.mark.skipif(sys.platform == "darwin", reason="bug with Mac with pygame")
+@pytest.mark.parametrize("rendering_tool", RENDERING_TOOL)
+def test_gym_make_rendering_gif(rendering_tool):
+    seeder = Seeder(123)
+    env = gym_make("MountainCar-v0", render_mode="rgb_array")
+    env.reseed(seeder)
+    env.renderer_type = rendering_tool
+
+    agent = RSUCBVIAgent(
+        env,
+        gamma=0.99,
+        horizon=200,
+        bonus_scale_factor=0.1,
+        copy_env=False,
+        min_dist=0.1,
+    )
+
+    info = agent.fit(15)
+    print(info)
+
+    env.enable_rendering()
+    observation, info = env.reset()
+    for tt in range(100):
+        action = agent.policy(observation)
+        observation, reward, terminated, truncated, info = env.step(action)
+        done = terminated or truncated
+        if done:
+            # Warning: this will never happen in the present case.
+            break
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        saving_path = tmpdirname + "/test_gif.gif"
+        env.save_gif(saving_path)
+        assert os.path.isfile(saving_path)
+
+        try:
+            os.remove(saving_path)
+        except Exception:
+            pass
+
+
+# ##### Works only if you add a dependency to ffmpeg ############
+# @pytest.mark.skipif(sys.platform == "darwin", reason="bug with Mac with pygame")
+# @pytest.mark.parametrize("rendering_tool", RENDERING_TOOL)
+# def test_gym_make_rendering_mp4(rendering_tool):
+#     seeder = Seeder(123)
+#     env = gym_make("MountainCar-v0", render_mode="rgb_array")
+#     env.reseed(seeder)
+#     env.renderer_type = rendering_tool
+
+#     agent = RSUCBVIAgent(
+#         env,
+#         gamma=0.99,
+#         horizon=200,
+#         bonus_scale_factor=0.1,
+#         copy_env=False,
+#         min_dist=0.1,
+#     )
+
+#     info = agent.fit(15)
+#     print(info)
+
+#     env.enable_rendering()
+#     observation, info = env.reset()
+#     for tt in range(100):
+#         action = agent.policy(observation)
+#         observation, reward, terminated, truncated, info = env.step(action)
+#         done = terminated or truncated
+#         if done:
+#             # Warning: this will never happen in the present case.
+#             break
+
+#     with tempfile.TemporaryDirectory() as tmpdirname:
+#         saving_path = tmpdirname + "/test_video.mp4"
+#         env.save_video(saving_path)
+#         assert os.path.isfile(saving_path)
+#         try:
+#             os.remove(saving_path)
+#         except Exception:
+#             pass