code: def make_networks(
      spec: specs.EnvironmentSpec, discrete_actions: bool = False
  ) -> bc.BCNetworks:
      """Creates networks used by the agent."""
  
      if discrete_actions:
          final_layer_size = spec.actions.num_values
      else:
          final_layer_size = np.prod(spec.actions.shape, dtype=int)
  
      def _actor_fn(obs, is_training=False, key=None):
          # is_training and key allows to defined train/test dependant modules
          # like dropout.
          del is_training
          del key
          if discrete_actions:
              network = hk.nets.MLP([64, 64, final_layer_size])
          else:
              network = hk.Sequential(
                  [
                      networks_lib.LayerNormMLP(
                          [64, 64], activate_final=True
                      ),
                      networks_lib.NormalTanhDistribution(
                          final_layer_size
                      ),
                  ]
              )
          return network(obs)
  
      policy = hk.without_apply_rng(hk.transform(_actor_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)
      policy_network = networks_lib.FeedForwardNetwork(
          lambda key: policy.init(key, dummy_obs), policy.apply
      )
      bc_policy_network = bc.convert_to_bc_network(policy_network)
  
      if discrete_actions:
  
          def sample_fn(
              logits: networks_lib.NetworkOutput, key: jax_types.PRNGKey
          ) -> networks_lib.Action:
              return rlax.epsilon_greedy(epsilon=0.0).sample(
                  key, logits
              )
  
          def log_prob(
              logits: networks_lib.NetworkOutput,
              actions: networks_lib.Action,
          ) -> networks_lib.LogProb:
              max_logits = jnp.max(logits, axis=-1, keepdims=True)
              logits = logits - max_logits
              logits_actions = jnp.sum(
                  jax.nn.one_hot(actions, spec.actions.num_values)
                  * logits,
                  axis=-1,
              )
  
              log_prob = logits_actions - special.logsumexp(
                  logits, axis=-1
              )
              return log_prob
  
      else:
  
          def sample_fn(
              distribution: networks_lib.NetworkOutput,
              key: jax_types.PRNGKey,
          ) -> networks_lib.Action:
              return distribution.sample(seed=key)
  
          def log_prob(
              distribuition: networks_lib.NetworkOutput,
              actions: networks_lib.Action,
          ) -> networks_lib.LogProb:
              return distribuition.log_prob(actions)
  
      return bc.BCNetworks(bc_policy_network, sample_fn, log_prob)
  
imports: from acme import specs
  from acme import types
  from acme.agents.jax import bc
  from acme.jax import networks
  from acme.jax import types
  from acme.jax import utils
  from acme.testing import fakes
  import chex
  import haiku
  import jax
  import jax.numpy
  from jax.scipy import special
  import numpy
  import optax
  import rlax
  from absl.testing import absltest
  from absl.testing import parameterized
label_desc: Use image node_id 1 for calling a global function with example usage: make_networks(spec, discrete_actions) and returns: bc, network, rlax, log_prob, distribution, distribuition