import tensorflow as tf

  def f(layer_input):
    return tf.keras.layers.concatenate(list(layer_input), axis=-1)
  return tf.keras.layers.Lambda(f)


@gin.configurable
def extract_observation_layer():
  return tf.keras.layers.Lambda(lambda obs: obs['observation'])


@gin.configurable
def monitor_freq(freq=100, vid_dir='./videos'):
  return functools.partial(Monitor, video_callable=lambda x: (x%freq) == 0,
                           directory=vid_dir)
