Check if ffmpeg is installed and guard against it in RL env.

PiperOrigin-RevId: 228910916

Author: Lukasz Kaiser

tensor2tensor/layers/common_video.py

@@ -357,8 +357,8 @@ def _encode_gif(images, fps):
   """Encodes numpy images into gif string.
 
   Args:
-    images: A 5-D `uint8` `np.array` (or a list of 4-D images) of shape
-      `[batch_size, time, height, width, channels]` where `channels` is 1 or 3.
+    images: A 4-D `uint8` `np.array` (or a list of 3-D images) of shape
+      `[time, height, width, channels]` where `channels` is 1 or 3.
     fps: frames per second of the animation
 
   Returns:
@@ -372,6 +372,16 @@ def _encode_gif(images, fps):
   return writer.finish()
 
 
+def ffmpeg_works():
+  """Tries to encode images with ffmpeg to check if it works."""
+  images = np.zeros((2, 32, 32, 3), dtype=np.uint8)
+  try:
+    _encode_gif(images, 2)
+    return True
+  except (IOError, OSError):
+    return False
+
+
 def py_gif_summary(tag, images, max_outputs, fps, return_summary_value=False):
   """Outputs a `Summary` protocol buffer with gif animations.
 
@@ -697,7 +707,7 @@ def __init__(self, fps, output_path=None, file_format="gif"):
   def __init_ffmpeg(self, image_shape):
     """Initializes ffmpeg to write frames."""
     import itertools  # pylint: disable=g-import-not-at-top
-    from subprocess import Popen, PIPE  # pylint: disable=g-import-not-at-top,g-multiple-import
+    from subprocess import Popen, PIPE  # pylint: disable=g-import-not-at-top,g-multiple-import,g-importing-member
     ffmpeg = "ffmpeg"
     height, width, channels = image_shape
     self.cmd = [

tensor2tensor/rl/envs/simulated_batch_env.py

@@ -115,6 +115,7 @@ def __init__(
     """Batch of environments inside the TensorFlow graph."""
     super(SimulatedBatchEnv, self).__init__(observation_space, action_space)
 
+    self._ffmpeg_works = common_video.ffmpeg_works()
     self.batch_size = batch_size
     self._min_reward = reward_range[0]
     self._num_frames = frame_stack_size
@@ -267,6 +268,8 @@ def history_observations(self):
     return self.history_buffer.get_all_elements()
 
   def _video_dump_frame(self, obs, rews):
+    if not self._ffmpeg_works:
+      return
     if self._video_writer is None:
       self._video_counter += 1
       self._video_writer = common_video.WholeVideoWriter(
@@ -280,6 +283,8 @@ def _video_dump_frame(self, obs, rews):
     self._video_writer.write(np.concatenate([np.asarray(img), obs[0]], axis=0))
 
   def _video_dump_frames(self, obs):
+    if not self._ffmpeg_works:
+      return
     zeros = np.zeros(obs.shape[0])
     for i in range(obs.shape[1]):
       self._video_dump_frame(obs[:, i, :], zeros)