code: def policy_prior_loss(
      apply_fn: Callable[
          [networks.Observation, networks.Action], networks.Action
      ],
      steps: types.Transition,
  ):
      """Returns the loss for the policy prior.
  
      Args:
        apply_fn: applies a policy prior (o_t, a_t) -> a_t+1, 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_tm1 = steps.action[:, dataset.PREVIOUS, ...]
      action_t = steps.action[:, dataset.CURRENT, ...]
  
      predicted_action_t = apply_fn(observation_t, action_tm1)
      return mse(predicted_action_t, action_t)
  
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 3 for calling a global function with example usage: policy_prior_loss(apply_fn, steps) and returns: mse