Merge remote-tracking branch 'refs/remotes/origin/fix-codec' into fix-codec

Author: MGross21

mujoco_toolbox/sim.py

@@ -17,14 +17,14 @@
 from typing import TYPE_CHECKING, Any, Self, TypeAlias
 
 import defusedxml.ElementTree as ET
-import matplotlib.pyplot as plt
-import matplotlib.animation as animation
 import imageio.v3 as iio
+import matplotlib.pyplot as plt
 import mujoco
 import mujoco.viewer
 import numpy as np
 import yaml
-from IPython.display import clear_output, HTML
+from IPython.display import HTML, clear_output
+from matplotlib import animation
 from tqdm.auto import tqdm
 
 from .builder import Builder
@@ -715,7 +715,7 @@ def is_jupyter() -> bool:
 
             # Set up the figure and image once
             fig, ax = plt.subplots()
-            im = ax.imshow(np.zeros((self.resolution[1], self.resolution[0], 3), dtype=np.uint8), interpolation='nearest')
+            im = ax.imshow(np.zeros((self.resolution[1], self.resolution[0], 3), dtype=np.uint8), interpolation="nearest")
             ax.set_axis_off()
             ax.set_title(title)
             plt.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=0, hspace=0)
@@ -726,18 +726,17 @@ def is_jupyter() -> bool:
                     lambda frame: (im.set_data(frame), [im])[1],
                     frames=subset_frames,
                     interval=(1000 / self._fps),
-                    blit=True
+                    blit=True,
                 )
                 plt.close(fig)
                 return HTML(ani.to_jshtml())
-            else:
-                plt.ion()
-                delay = 1.0 / self._fps
-                for frame in subset_frames:
-                    im.set_data(frame)
-                    plt.pause(delay)
-                plt.ioff()
-                plt.close(fig)
+            plt.ion()
+            delay = 1.0 / self._fps
+            for frame in subset_frames:
+                im.set_data(frame)
+                plt.pause(delay)
+            plt.ioff()
+            plt.close(fig)
         except Exception as e:
             msg = "Error while showing video subset."
             raise Exception(msg) from e  # noqa: TRY002
@@ -769,15 +768,15 @@ def save(
         if not hasattr(self, "_frames") or self._frames is None or len(self._frames) == 0:
             msg = "No frames captured to render. Re-run the simulation with render=True."
             raise ValueError(msg)
-        
+
         # Extract frames
         subset_frames = self._get_index(
             frame_idx=frame_idx,
             time_idx=time_idx,
         )
 
         title_path = Path(title)
-    
+
         try:
             # Save the video
             iio.imwrite(

tests/test_id.py

@@ -1,4 +1,3 @@
-import os
 
 import pytest
 

tests/test_media.py

@@ -1,7 +1,9 @@
-import mujoco_toolbox as mjtb
-from pathlib import Path
-import pytest
 import os
+from pathlib import Path
+
+import mujoco_toolbox as mjtb
+
+FILE_EXTENSIONS = ["mp4", "gif", "avi", "webm", "mov"]
 
 FILE_EXTENSIONS = ["mp4", "gif", "avi", "webm", "mov"]
 

tests/test_utils.py

@@ -1,4 +1,3 @@
-import os
 
 from mujoco_toolbox import utils
 

tests/test_warnings.py

@@ -1,28 +1,30 @@
-from mujoco_toolbox.warnings import SimulationError, SimulationWarning
-from mujoco_toolbox.sim import Wrapper
-import mujoco_toolbox as mjtb
 import pytest
 
+import mujoco_toolbox as mjtb
+from mujoco_toolbox.sim import Wrapper
+from mujoco_toolbox.warnings import SimulationError, SimulationWarning
+
+
 # Needs warnings to be enabled (Dummy test to ensure warnings are raised)
-def test_simulation_warning():
+def test_simulation_warning() -> None:
     msg = "This is a test warning"
     try:
         raise SimulationWarning(msg)
     except SimulationWarning as e:
         assert str(e) == msg
 
-def test_simulation_error():
+def test_simulation_error() -> None:
     msg = "This is a test error"
     try:
         raise SimulationError(msg)
     except SimulationError as e:
         assert str(e) == msg
 
-def test_deprecated_wrapper_warning():
+def test_deprecated_wrapper_warning() -> None:
     with pytest.warns(DeprecationWarning):
         Wrapper("<mujoco></mujoco>")
-    
-def test_deprecated_wrapper_error():
+
+def test_deprecated_wrapper_error() -> None:
     mjtb.__version__ = "1.0.0"  # Set version to trigger deprecation
     with pytest.raises(RuntimeError):
-        Wrapper("<mujoco></mujoco>")
+        Wrapper("<mujoco></mujoco>")