code: def compute_rnd_reward(
      predictor_params: networks_lib.Params,
      target_params: networks_lib.Params,
      transitions: types.Transition,
      networks: RNDNetworks,
  ) -> jnp.ndarray:
      """Computes the intrinsic RND reward for a given transition.
  
      Args:
        predictor_params: Parameters of the predictor network.
        target_params: Parameters of the target network.
        transitions: The sample to compute rewards for.
        networks: RND networks
  
      Returns:
        The rewards as an ndarray.
      """
      target_output = networks.target.apply(
          target_params, transitions.observation, transitions.action
      )
      predictor_output = networks.predictor.apply(
          predictor_params, transitions.observation, transitions.action
      )
      return networks.get_reward(
          predictor_output, target_output, transitions.reward
      )
  
file_path: acme/acme/agents/jax/rnd/networks.py
imports: import dataclasses
  import functools
  from typing import Callable, Generic, Tuple, TypeVar
  from acme import specs
  from acme import types
  from acme.jax import networks
  from acme.jax import utils
  import haiku
  import jax.numpy
label_desc: Use image node_id 3 for calling a global function with example usage: compute_rnd_reward(predictor_params, target_params, transitions, networks) and returns: networks