code: def make_n_step_return_network(
      spec: specs.EnvironmentSpec,
      hidden_layer_sizes: Tuple[int, ...] = (64, 64),
  ) -> networks.FeedForwardNetwork:
      """Creates an N-step return network used by the agent."""
  
      def _n_step_return_fn(
          observation_t, action_t, is_training=False, key=None
      ):
          # is_training and key allows to defined train/test dependant modules
          # like dropout.
          del is_training
          del key
          network = hk.nets.MLP(hidden_layer_sizes + (1,))
          return network(
              jnp.concatenate([observation_t, action_t], axis=-1)
          )
  
      n_step_return = hk.without_apply_rng(
          hk.transform(_n_step_return_fn)
      )
      return make_network_from_module(n_step_return, spec)
  
file_path: acme/acme/agents/jax/mbop/networks.py
imports: import dataclasses
  from typing import Any, Tuple
  from acme import specs
  from acme.jax import networks
  from acme.jax import utils
  import haiku
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: make_n_step_return_network(spec, hidden_layer_sizes) and returns: make_network_from_module, network