code: def make_networks(
      spec: specs.EnvironmentSpec,
      hidden_layer_sizes: Sequence[int] = (256, 256),
  ) -> TD3Networks:
      """Creates networks used by the agent.
  
      The networks used are based on LayerNormMLP, which is different than the
      MLP with relu activation described in TD3 (which empirically performs worse).
  
      Args:
        spec: Environment specs
        hidden_layer_sizes: list of sizes of hidden layers in actor/critic networks
  
      Returns:
        network: TD3Networks
      """
  
      action_specs = spec.actions
      num_dimensions = np.prod(action_specs.shape, dtype=int)
  
      def add_policy_noise(
          action: types.NestedArray,
          key: networks_lib.PRNGKey,
          target_sigma: float,
          noise_clip: float,
      ) -> types.NestedArray:
          """Adds action noise to bootstrapped Q-value estimate in critic loss."""
          noise = (
              jax.random.normal(key=key, shape=action_specs.shape)
              * target_sigma
          )
          noise = jnp.clip(noise, -noise_clip, noise_clip)
          return jnp.clip(
              action + noise, action_specs.minimum, action_specs.maximum
          )
  
      def _actor_fn(obs: types.NestedArray) -> types.NestedArray:
          network = hk.Sequential(
              [
                  networks_lib.LayerNormMLP(
                      hidden_layer_sizes, activate_final=True
                  ),
                  networks_lib.NearZeroInitializedLinear(
                      num_dimensions
                  ),
                  networks_lib.TanhToSpec(spec.actions),
              ]
          )
          return network(obs)
  
      def _critic_fn(
          obs: types.NestedArray, action: types.NestedArray
      ) -> types.NestedArray:
          network1 = hk.Sequential(
              [
                  networks_lib.LayerNormMLP(
                      list(hidden_layer_sizes) + [1]
                  ),
              ]
          )
          input_ = jnp.concatenate([obs, action], axis=-1)
          value = network1(input_)
          return jnp.squeeze(value)
  
      policy = hk.without_apply_rng(hk.transform(_actor_fn))
      critic = hk.without_apply_rng(hk.transform(_critic_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)
  
      network = TD3Networks(
          policy_network=networks_lib.FeedForwardNetwork(
              lambda key: policy.init(key, dummy_obs), policy.apply
          ),
          critic_network=networks_lib.FeedForwardNetwork(
              lambda key: critic.init(key, dummy_obs, dummy_action),
              critic.apply,
          ),
          twin_critic_network=networks_lib.FeedForwardNetwork(
              lambda key: critic.init(key, dummy_obs, dummy_action),
              critic.apply,
          ),
          add_policy_noise=add_policy_noise,
      )
  
      return network
  
file_path: acme/acme/agents/jax/td3/networks.py
imports: import dataclasses
  from typing import Callable, Sequence
  from acme import specs
  from acme import types
  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: network, jnp, network, jnp