code: def world_model_loss(
      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),
      )
  
file_path: acme/acme/agents/jax/mbop/losses.py
imports: import dataclasses
  from typing import Any, Callable, Optional, Tuple, Union
  from acme import types
  from acme.agents.jax.mbop import dataset
  from acme.jax import networks
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: world_model_loss(apply_fn, steps) and returns: mse