code: def make_continuous_networks(
      environment_spec: specs.EnvironmentSpec,
      policy_layer_sizes: Sequence[int] = (64, 64),
      value_layer_sizes: Sequence[int] = (64, 64),
      use_tanh_gaussian_policy: bool = True,
  ) -> PPONetworks:
      """Creates PPONetworks to be used for continuous action environments."""
  
      # Get total number of action dimensions from action spec.
      num_dimensions = np.prod(
          environment_spec.actions.shape, dtype=int
      )
  
      def forward_fn(inputs: networks_lib.Observation):
          def _policy_network(obs: networks_lib.Observation):
              h = utils.batch_concat(obs)
              h = hk.nets.MLP(policy_layer_sizes, activate_final=True)(
                  h
              )
  
              # tfd distributions have a weird bug in jax when vmapping is used, so the
              # safer implementation in general is for the policy network to output the
              # distribution parameters, and for the distribution to be constructed
              # in a method such as make_ppo_networks above
              if not use_tanh_gaussian_policy:
                  # Following networks_lib.MultivariateNormalDiagHead
                  init_scale = 0.3
                  min_scale = 1e-6
                  w_init = hk.initializers.VarianceScaling(1e-4)
                  b_init = hk.initializers.Constant(0.0)
                  loc_layer = hk.Linear(
                      num_dimensions, w_init=w_init, b_init=b_init
                  )
                  scale_layer = hk.Linear(
                      num_dimensions, w_init=w_init, b_init=b_init
                  )
  
                  loc = loc_layer(h)
                  scale = jax.nn.softplus(scale_layer(h))
                  scale *= init_scale / jax.nn.softplus(0.0)
                  scale += min_scale
  
                  return MVNDiagParams(loc=loc, scale_diag=scale)
  
              # Following networks_lib.NormalTanhDistribution
              min_scale = 1e-3
              w_init = hk.initializers.VarianceScaling(
                  1.0, "fan_in", "uniform"
              )
              b_init = hk.initializers.Constant(0.0)
              loc_layer = hk.Linear(
                  num_dimensions, w_init=w_init, b_init=b_init
              )
              scale_layer = hk.Linear(
                  num_dimensions, w_init=w_init, b_init=b_init
              )
  
              loc = loc_layer(h)
              scale = scale_layer(h)
              scale = jax.nn.softplus(scale) + min_scale
  
              return TanhNormalParams(loc=loc, scale=scale)
  
          value_network = hk.Sequential(
              [
                  utils.batch_concat,
                  hk.nets.MLP(value_layer_sizes, activate_final=True),
                  hk.Linear(1),
                  lambda x: jnp.squeeze(x, axis=-1),
              ]
          )
  
          policy_output = _policy_network(inputs)
          value = value_network(inputs)
          return (policy_output, value)
  
      # Transform into pure functions.
      forward_fn = hk.without_apply_rng(hk.transform(forward_fn))
  
      dummy_obs = utils.zeros_like(environment_spec.observations)
      dummy_obs = utils.add_batch_dim(
          dummy_obs
      )  # Dummy 'sequence' dim.
      network = networks_lib.FeedForwardNetwork(
          lambda rng: forward_fn.init(rng, dummy_obs), forward_fn.apply
      )
  
      # Create PPONetworks to add functionality required by the agent.
  
      if not use_tanh_gaussian_policy:
          return make_mvn_diag_ppo_networks(network)
  
      return make_tanh_normal_ppo_networks(network)
  
imports: import dataclasses
  from typing import Callable, NamedTuple, Optional, Sequence
  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
  import tensorflow_probability
label_desc: Use image node_id 8 for calling a global function with example usage: make_continuous_networks(environment_spec, policy_layer_sizes, value_layer_sizes, use_tanh_gaussian_policy) and returns: make_tanh_normal_ppo_networks, policy_output,  value, make_mvn_diag_ppo_networks, TanhNormalParams, MVNDiagParams