import tensorflow as tf

    num_target_frames = 1
  target_value_shape_suffix = [num_target_frames]
  if distributional_size > 1:
    target_value_shape_suffix = [num_target_frames, distributional_size]
  features = {
      "inputs": observations,
      "epoch": tf.constant(epoch + 1),
      "input_action": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
      "input_reward": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
      "targets": tf.zeros(obs_shape[:1] + [num_target_frames] + obs_shape[2:]),
      "target_action": tf.zeros(
          obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
