Transforms time step observations to frames of a video.

def _generate_time_steps(self, trajectory_list):
    """Transforms time step observations to frames of a video."""
    for time_step in env_problem.EnvProblem._generate_time_steps(
        self, trajectory_list):
      # Convert the rendered observations from numpy to png format.
      frame_np = np.array(time_step.pop(env_problem.OBSERVATION_FIELD))
      frame_np = frame_np.reshape(
          [self.frame_height, self.frame_width, self.num_channels])
      # TODO(msaffar) Add support for non RGB rendered environments
      frame = png.from_array(frame_np, "RGB", info={"bitdepth": 8})
      frame_buffer = six.BytesIO()
      frame.save(frame_buffer)

      # Put the encoded frame back.
      time_step[_IMAGE_ENCODED_FIELD] = [frame_buffer.getvalue()]
      time_step[_IMAGE_FORMAT_FIELD] = [_FORMAT]
      time_step[_IMAGE_HEIGHT_FIELD] = [self.frame_height]
      time_step[_IMAGE_WIDTH_FIELD] = [self.frame_width]

      # Add the frame number
      time_step[_FRAME_NUMBER_FIELD] = time_step[env_problem.TIMESTEP_FIELD]
      yield time_step