code: def make_networks(
      spec: specs.EnvironmentSpec,
      hidden_layer_sizes: Tuple[int, ...] = (256, 256),
  ) -> ValueDiceNetworks:
      """Creates networks used by the agent."""
  
      num_dimensions = np.prod(spec.actions.shape, dtype=int)
  
      def _actor_fn(obs):
          network = hk.Sequential(
              [
                  hk.nets.MLP(
                      list(hidden_layer_sizes),
                      w_init=hk.initializers.VarianceScaling(
                          1.0, "fan_in", "uniform"
                      ),
                      activation=jax.nn.relu,
                      activate_final=True,
                  ),
                  networks_lib.NormalTanhDistribution(num_dimensions),
              ]
          )
          return network(obs)
  
      def _nu_fn(obs, action):
          network = hk.Sequential(
              [
                  hk.nets.MLP(
                      list(hidden_layer_sizes) + [1],
                      w_init=hk.initializers.VarianceScaling(
                          1.0, "fan_in", "uniform"
                      ),
                      activation=jax.nn.relu,
                  ),
              ]
          )
          return network(jnp.concatenate([obs, action], axis=-1))
  
      policy = hk.without_apply_rng(hk.transform(_actor_fn))
      nu = hk.without_apply_rng(hk.transform(_nu_fn))
  
      # Create dummy observations and actions to create network parameters.
      dummy_action = utils.zeros_like(spec.actions)
      dummy_obs = utils.zeros_like(spec.observations)
      dummy_action = utils.add_batch_dim(dummy_action)
      dummy_obs = utils.add_batch_dim(dummy_obs)
  
      return ValueDiceNetworks(
          policy_network=networks_lib.FeedForwardNetwork(
              lambda key: policy.init(key, dummy_obs), policy.apply
          ),
          nu_network=networks_lib.FeedForwardNetwork(
              lambda key: nu.init(key, dummy_obs, dummy_action),
              nu.apply,
          ),
          sample=lambda params, key: params.sample(seed=key),
          sample_eval=lambda params, key: params.mode(),
      )
  
file_path: acme/acme/agents/jax/value_dice/networks.py
imports: import dataclasses
  from typing import Callable, Optional, Tuple
  from acme import specs
  from acme.agents.jax import actor_core
  from acme.jax import networks
  from acme.jax import utils
  import haiku
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: make_networks(spec, hidden_layer_sizes) and returns: ValueDiceNetworks, network, network