def METHOD_NAME(
    state: RNDTrainingState, transitions: types.Transition,
    loss_fn: RNDLoss, optimizer: optax.GradientTransformation
) -> Tuple[RNDTrainingState, Dict[str, jnp.ndarray]]:
  """Run an update steps on the given transitions.
  Args:
    state: The learner state.
    transitions: Transitions to update on.
    loss_fn: The loss function.
    optimizer: The optimizer of the predictor network.
  Returns:
    A new state and metrics.
  """
  loss, grads = jax.value_and_grad(loss_fn)(
      state.params,
      state.target_params,
      transitions=transitions)
  update, optimizer_state = optimizer.update(grads, state.optimizer_state)
  params = optax.apply_updates(state.params, update)
  new_state = RNDTrainingState(
      optimizer_state=optimizer_state,
      params=params,
      target_params=state.target_params,
      steps=state.steps + 1,
  )
  return new_state, {'rnd_loss': loss}