import tensorflow as tf

      next_states, next_tasks = self._task_distribution.split(
          experience.observation[:, 1])
      rewards, dones = self._task_distribution.evaluate(states,
                                                        experience.action[:, 0],
                                                        tasks)
      # Strictly speaking, we don't need to relabel the next rewards and next
      # dones because they end up being thrown away. Only the current rewards
      # and dones end up being important.
      next_rewards, next_dones = self._task_distribution.evaluate(
          next_states, experience.action[:, 1], next_tasks)

      new_rewards = tf.concat([rewards[:, None], next_rewards[:, None]], axis=1)
      new_dones = tf.concat([dones[:, None], next_dones[:, None]], axis=1)
      # 0 if episode is done, 1 if episode is continuing
      new_discount = 1.0 - tf.cast(new_dones, tf.float32)
      assert new_rewards.shape == experience.reward.shape
      assert new_discount.shape == experience.discount.shape
      experience = experience.replace(reward=new_rewards, discount=new_discount)
    return experience

  def _soft_relabel(self, experience):
    """Reassigns tasks to each state and next state.

    Does not recompute the rewards or done flags.

    Args:
      experience: The experience that we want to relabel with inverse RL.
    Returns:
