add docstring the gif_recorder.py and recorder_wrapper.py.

setup.py

@@ -69,6 +69,7 @@
             "sphinxcontrib-apidoc>=0.3.0",
         ],
         "extras": ["pynput>=1.7.4"],  # Used by HumanKeyboardAgent
+        "gym": ["moviepy = 1.0.3"]
         "remote_agent": ["grpcio==1.32.0"],
         "rllib": [
             "opencv-python==4.1.2.30",

smarts/env/wrappers/gif_recorder.py

@@ -21,15 +21,18 @@
 # THE SOFTWARE.
 import os
 
-from moviepy.editor import *
-import gym.envs
+from moviepy.editor import ImageClip
+from moviepy.editor import ImageSequenceClip
 import shutil
 import time
 from pathlib import Path
 
 
 class GifRecorder:
     def __init__(self, dir, env):
+        """
+        Use images(rgb_array) to create a gif file.
+        """
         timestamp_str = time.strftime("%Y%m%d-%H%M%S")
         self.dir = dir + "_" + timestamp_str
         self.env = env
@@ -42,15 +45,24 @@ def __init__(self, dir, env):
         self._dir_name = str(Path(dir).name)
 
     def capture_frame(self, step_num, image):
+        """
+        Create image according to the rgb_array and store it with step number in the destinated folder
+        """
         with ImageClip(image) as image_clip:
             image_clip.save_frame(f"{self.dir}/{self._dir_name}_{step_num}.jpeg")
 
     def generate_gif(self):
+        """
+        Use the images in the same folder to create a gif file.
+        """
         with ImageSequenceClip(self.dir, fps=10) as clip:
             clip.write_gif(f"videos/{self._dir_name}.gif")
         clip.close()
 
     def close_recorder(self):
+        """
+        close the recorder by deleting the image folder.
+        """
         try:
             shutil.rmtree(self.dir)
         except:

smarts/env/wrappers/recorder_wrapper.py

@@ -29,6 +29,9 @@
 
 
 class RecorderWrapper(gym.Wrapper):
+    """
+    A Wrapper that interacts the gym environment with the GifRecorder to record video step by step.
+    """
     def __init__(self, dir, env):
 
         try:
@@ -44,23 +47,35 @@ def __init__(self, dir, env):
         self.current_frame = -1
 
     def reset(self, **kwargs):
+        """
+        Reset the gym environment and restart recording.
+        """
         observations = super().reset(**kwargs)
         if self.recording == False:
             self.start_recording()
 
         return observations
 
     def start_recording(self):
+        """
+        Start the gif recorder and capture the first frame.
+        """
         if self.gif_recorder is None:
             self.gif_recorder = GifRecorder(self.dir, self.env)
         image = super().render(mode="rgb_array")
         self.gif_recorder.capture_frame(self.next_frame_id(),image)
         self.recording = True
 
     def stop_recording(self):
+        """
+        Stop recording.
+        """
         self.recording = False
 
     def step(self, action):
+        """
+        Step the environment using the action and record the next frame.
+        """
         observations, rewards, dones, infos = super().step(action)
         if self.recording == True:
             image = super().render(mode="rgb_array")
@@ -69,10 +84,16 @@ def step(self, action):
         return observations, rewards, dones, infos
 
     def next_frame_id(self):
+        """
+        Get the id for next frame.
+        """
         self.current_frame += 1
         return self.current_frame
 
     def close(self):
+        """
+        Close the recorder by deleting the image folder and generate the gif file.
+        """
         if self.gif_recorder is not None:
             self.gif_recorder.generate_gif()
             self.gif_recorder.close_recorder()