code: def make_networks(
      spec: specs.EnvironmentSpec,
      direct_rl_networks: DirectRLNetworks,
      layer_sizes: Tuple[int, ...] = (256, 256),
      intrinsic_reward_coefficient: float = 1.0,
      extrinsic_reward_coefficient: float = 0.0,
  ) -> RNDNetworks[DirectRLNetworks]:
      """Creates networks used by the agent and returns RNDNetworks.
  
      Args:
        spec: Environment spec.
        direct_rl_networks: Networks used by a direct rl algorithm.
        layer_sizes: Layer sizes.
        intrinsic_reward_coefficient: Multiplier on intrinsic reward.
        extrinsic_reward_coefficient: Multiplier on extrinsic reward.
  
      Returns:
        The RND networks.
      """
  
      def _rnd_fn(obs, act):
          # RND does not use the action but other variants like RED do.
          del act
          network = networks_lib.LayerNormMLP(list(layer_sizes))
          return network(obs)
  
      target = hk.without_apply_rng(hk.transform(_rnd_fn))
      predictor = hk.without_apply_rng(hk.transform(_rnd_fn))
  
      # Create dummy observations and actions to create network parameters.
      dummy_obs = utils.zeros_like(spec.observations)
      dummy_obs = utils.add_batch_dim(dummy_obs)
  
      return RNDNetworks(
          target=networks_lib.FeedForwardNetwork(
              lambda key: target.init(key, dummy_obs, ()), target.apply
          ),
          predictor=networks_lib.FeedForwardNetwork(
              lambda key: predictor.init(key, dummy_obs, ()),
              predictor.apply,
          ),
          direct_rl_networks=direct_rl_networks,
          get_reward=functools.partial(
              rnd_reward_fn,
              intrinsic_reward_coefficient=intrinsic_reward_coefficient,
              extrinsic_reward_coefficient=extrinsic_reward_coefficient,
          ),
      )
  
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 2 for calling a global function with example usage: make_networks(spec, direct_rl_networks, layer_sizes, intrinsic_reward_coefficient, extrinsic_reward_coefficient) and returns: RNDNetworks, network