code: def make_network_with_prior(
      action_spec: specs.BoundedArray,
      policy_layer_sizes: Sequence[int] = (200, 100),
      critic_layer_sizes: Sequence[int] = (400, 300),
      prior_layer_sizes: Sequence[int] = (200, 100),
      policy_keys: Optional[Sequence[str]] = None,
      prior_keys: Optional[Sequence[str]] = None,
  ) -> Mapping[str, types.TensorTransformation]:
      """Creates networks used by the agent."""
  
      # Get total number of action dimensions from action spec.
      num_dimensions = np.prod(action_spec.shape, dtype=int)
      flatten_concat_policy = functools.partial(
          svg0_utils.batch_concat_selection, concat_keys=policy_keys
      )
      flatten_concat_prior = functools.partial(
          svg0_utils.batch_concat_selection, concat_keys=prior_keys
      )
  
      policy_network = snt.Sequential(
          [
              flatten_concat_policy,
              networks.LayerNormMLP(
                  policy_layer_sizes, activate_final=True
              ),
              networks.MultivariateNormalDiagHead(
                  num_dimensions,
                  tanh_mean=True,
                  min_scale=0.1,
                  init_scale=0.7,
                  fixed_scale=False,
                  use_tfd_independent=False,
              ),
          ]
      )
      # The multiplexer concatenates the (maybe transformed) observations/actions.
      multiplexer = networks.CriticMultiplexer(
          observation_network=flatten_concat_policy,
          action_network=networks.ClipToSpec(action_spec),
      )
      critic_network = snt.Sequential(
          [
              multiplexer,
              networks.LayerNormMLP(
                  critic_layer_sizes, activate_final=True
              ),
              networks.NearZeroInitializedLinear(1),
          ]
      )
      prior_network = snt.Sequential(
          [
              flatten_concat_prior,
              networks.LayerNormMLP(
                  prior_layer_sizes, activate_final=True
              ),
              networks.MultivariateNormalDiagHead(
                  num_dimensions,
                  tanh_mean=True,
                  min_scale=0.1,
                  init_scale=0.7,
                  fixed_scale=False,
                  use_tfd_independent=False,
              ),
          ]
      )
      return {
          "policy": policy_network,
          "critic": critic_network,
          "prior": prior_network,
      }
  
file_path: acme/acme/agents/tf/svg0_prior/networks.py
imports: import functools
  from typing import Mapping, Sequence, Optional
  from acme import specs
  from acme import types
  from acme.agents.tf.svg0_prior import utils
  from acme.tf import networks
  from acme.tf import utils
  import numpy
  import sonnet
label_desc: Use image node_id 2 for calling a global function with example usage: make_network_with_prior(action_spec, policy_layer_sizes, critic_layer_sizes, prior_layer_sizes, policy_keys, prior_keys) and returns: dict