def METHOD_NAME(
    predictor_params: networks_lib.Params,
    target_params: networks_lib.Params,
    transitions: types.Transition,
    networks: rnd_networks.RNDNetworks,
) -> float:
  """The Random Network Distillation loss.
  See https://arxiv.org/pdf/1810.12894.pdf A.2
  Args:
    predictor_params: Parameters of the predictor
    target_params: Parameters of the target
    transitions: Transitions to compute the loss on.
    networks: RND networks
  Returns:
    The MSE loss as a float.
  """
  target_output = networks.target.apply(target_params,
                                        transitions.observation,
                                        transitions.action)
  predictor_output = networks.predictor.apply(predictor_params,
                                              transitions.observation,
                                              transitions.action)
  return jnp.mean(jnp.square(target_output - predictor_output))