//  [python]
def make_dataset_iterator(
      self, replay_client):
    """Create a dataset iterator to use for learning/updating the agent."""
    @tf.function
    def flatten_fn(sample):
      seq_len = tf.shape(sample.data.observation)[0]
      arange = tf.range(seq_len)
      is_future_mask = tf.cast(arange[:, None] < arange[None], tf.float32)
      discount = self._config.discount ** tf.cast(arange[None] - arange[:, None], tf.float32)  # pylint: disable=line-too-long
      probs = is_future_mask * discount
      # The indexing changes the shape from [seq_len, 1] to [seq_len]
      goal_index = tf.random.categorical(logits=tf.math.log(probs),
                                         num_samples=1)[:, 0]
      state = sample.data.observation[:-1, :self._config.obs_dim]
      next_state = sample.data.observation[1:, :self._config.obs_dim]

      # Create the goal observations in three steps.
      # 1. Take all future states (not future goals).
      # 2. Apply obs_to_goal.
      # 3. Sample one of the future states. Note that we don't look for a goal
      # for the final state, because there are no future states.
      goal = sample.data.observation[:, :self._config.obs_dim]
      goal = contrastive_utils.obs_to_goal_2d(
          goal, start_index=self._config.start_index,
          end_index=self._config.end_index)
      goal = tf.gather(goal, goal_index[:-1])
      new_obs = tf.concat([state, goal], axis=1)
      new_next_obs = tf.concat([next_state, goal], axis=1)
      transition = types.Transition(
          observation=new_obs,
          action=sample.data.action[:-1],
          reward=sample.data.reward[:-1],
          discount=sample.data.discount[:-1],
          next_observation=new_next_obs,
          extras={
              'next_action': sample.data.action[1:],
          })
      # Shift for the transpose_shuffle.
      shift = tf.random.uniform((), 0, seq_len, tf.int32)
      transition = tree.map_structure(lambda t: tf.roll(t, shift, axis=0),
                                      transition)
      return transition

    if self._config.num_parallel_calls:
      num_parallel_calls = self._config.num_parallel_calls
    else:
      num_parallel_calls = tf.data.AUTOTUNE

    def _make_dataset(unused_idx):
      dataset = reverb.TrajectoryDataset.from_table_signature(
          server_address=replay_client.server_address,
          table=self._config.replay_table_name,
          max_in_flight_samples_per_worker=100)
      dataset = dataset.map(flatten_fn)
      # transpose_shuffle
      def _transpose_fn(t):
        dims = tf.range(tf.shape(tf.shape(t))[0])
        perm = tf.concat([[1, 0], dims[2:]], axis=0)
        return tf.transpose(t, perm)
      dataset = dataset.batch(self._config.batch_size, drop_remainder=True)
      dataset = dataset.map(
          lambda transition: tree.map_structure(_transpose_fn, transition))
      dataset = dataset.unbatch()
      # end transpose_shuffle

      dataset = dataset.unbatch()
      return dataset
    dataset = tf.data.Dataset.from_tensors(0).repeat()
    dataset = dataset.interleave(
        map_func=_make_dataset,
        cycle_length=num_parallel_calls,
        num_parallel_calls=num_parallel_calls,
        deterministic=False)

    dataset = dataset.batch(
        self._config.batch_size * self._config.num_sgd_steps_per_step,
        drop_remainder=True)
    @tf.function
    def add_info_fn(data):
      info = reverb.SampleInfo(key=0,
                               probability=0.0,
                               table_size=0,
                               priority=0.0)
      return reverb.ReplaySample(info=info, data=data)
    dataset = dataset.map(add_info_fn, num_parallel_calls=tf.data.AUTOTUNE,
                          deterministic=False)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset.as_numpy_iterator()