def METHOD_NAME(apply_fn: Callable[[networks.Observation, networks.Action],
                                        Tuple[networks.Observation,
                                              jnp.ndarray]],
                     steps: types.Transition) -> jnp.ndarray:
  """Returns the loss for the world model.
  Args:
    apply_fn: applies a transition model (o_t, a_t) -> (o_t+1, r), expects the
      leading axis to index the batch and the second axis to index the
      transition triplet (t-1, t, t+1).
    steps: RLDS dictionary of transition triplets as prepared by
      `rlds_loader.episode_to_timestep_batch`.
  Returns:
    A scalar loss value as jnp.ndarray.
  """
  observation_t = jax.tree_map(lambda obs: obs[:, dataset.CURRENT, ...],
                               steps.observation)
  action_t = steps.action[:, dataset.CURRENT, ...]
  observation_tp1 = jax.tree_map(lambda obs: obs[:, dataset.NEXT, ...],
                                 steps.observation)
  reward_t = steps.reward[:, dataset.CURRENT, ...]
  (predicted_observation_tp1,
   predicted_reward_t) = apply_fn(observation_t, action_t)
  # predicted_* variables may have an extra outer dimension due to ensembling,
  # the mse loss still works due to broadcasting however.
  if len(observation_tp1.shape) != len(reward_t.shape):
    # The rewards from the transitions may not have the last singular dimension.
    reward_t = jnp.expand_dims(reward_t, axis=-1)
  return mse(
      jnp.concatenate([predicted_observation_tp1, predicted_reward_t], axis=-1),
      jnp.concatenate([observation_tp1, reward_t], axis=-1))