//  [python]
def update_step(self, replay_buffer_iter, train_target='both'):
    """Performs a single training step for critic and embedding.

    Args:
      replay_buffer_iter: A tensorflow graph iteratable object.
      train_target: string specifying whether update RL and or representation

    Returns:
      Dictionary with losses to track.
    """

    transition = next(replay_buffer_iter)
    numpy_dataset = isinstance(replay_buffer_iter, np.ndarray)
    # observation: n_batch x n_timesteps x 1 x H*W*3*n_frames x 1
    # -> n_batch x H x W x 3*n_frames
    if not numpy_dataset:
      states = transition.observation[:, 0]
      next_states = transition.observation[:, 1]
      actions = transition.action[:, 0]
      rewards = transition.reward
      level_ids = transition.policy_info[:, 0]
      if tf.shape(transition.reward)[1] > 2:
        rewards = tf.einsum(
            'ij,j->i', rewards,
            self.discount**tf.range(
                0, tf.shape(transition.reward)[1], dtype=tf.float32))
        self.n_step_rewards = tf.shape(transition.reward)[1]
      else:
        rewards = transition.reward[:, 0]
        self.n_step_rewards = 1
      discounts = transition.discount[:, 0]

      if transition.observation.dtype == tf.uint8:
        states = tf.cast(states, tf.float32) / 255.
        next_states = tf.cast(next_states, tf.float32) / 255.
    else:
      states, actions, rewards, next_states, discounts = transition

    self.reward_normalizer.update_normalization_statistics(rewards)

    if self.num_augmentations > 0:
      states, next_states = tf_utils.image_aug(
          states,
          next_states,
          img_pad=4,
          num_augmentations=self.num_augmentations,
          obs_dim=64,
          channels=3,
          cropped_shape=[self.batch_size, 68, 68, 3])

    next_actions_pi = self.act(next_states, data_aug=True)
    next_actions_mu = transition.action[:, 1]  # pylint: disable=unused-variable
    next_actions_pi_per_level = next_actions_mu

    states_b1 = states
    next_states_b1 = next_states
    actions_b1 = actions
    next_actions_b1 = next_actions_pi
    rewards_b1 = rewards
    discounts_b1 = discounts
    level_ids_b1 = level_ids

    states_b2 = states
    next_states_b2 = next_states
    actions_b2 = actions
    next_actions_b2 = next_actions_pi
    rewards_b2 = rewards
    discounts_b2 = discounts

    if train_target == 'encoder':
      print('Updating per-task critics')
      critic_distillation_dict = self.fit_task_critics(
          states_b1, actions_b1, next_states_b1, next_actions_pi_per_level,
          rewards_b1,
          discounts_b1, level_ids_b1)
      print('Done updating per-task critics')
      ssl_dict = {}
      critic_dict = {}
      return {**ssl_dict, **critic_distillation_dict}

    elif train_target == 'rl':
      critic_distillation_dict = {}
      print('Updating critic')
      critic_dict = self.fit_critic(states_b2, actions_b2, next_states_b2,
                                    next_actions_b2, rewards_b2, discounts_b2)
      print('Updating embedding')
      ssl_dict = self.fit_embedding(states_b1, actions_b1, next_states_b1,
                                    next_actions_b1, rewards_b1, discounts_b1,
                                    level_ids)
      print('Done')

    return {**ssl_dict, **critic_dict, **critic_distillation_dict}