code: def make_mvn_diag_ppo_networks(
      network: networks_lib.FeedForwardNetwork,
  ) -> PPONetworks:
      """Constructs a PPONetworks for MVN Diag policy from the FeedForwardNetwork.
  
      Args:
        network: a transformed Haiku network (or equivalent in other libraries) that
          takes in observations and returns the action distribution and value.
  
      Returns:
        A PPONetworks instance with pure functions wrapping the input network.
      """
  
      def log_prob(params: MVNDiagParams, action):
          return tfd.MultivariateNormalDiag(
              loc=params.loc, scale_diag=params.scale_diag
          ).log_prob(action)
  
      def entropy(
          params: MVNDiagParams, key: networks_lib.PRNGKey
      ) -> networks_lib.Entropy:
          del key
          return tfd.MultivariateNormalDiag(
              loc=params.loc, scale_diag=params.scale_diag
          ).entropy()
  
      def sample(params: MVNDiagParams, key: networks_lib.PRNGKey):
          return tfd.MultivariateNormalDiag(
              loc=params.loc, scale_diag=params.scale_diag
          ).sample(seed=key)
  
      def sample_eval(params: MVNDiagParams, key: networks_lib.PRNGKey):
          del key
          return tfd.MultivariateNormalDiag(
              loc=params.loc, scale_diag=params.scale_diag
          ).mode()
  
      return PPONetworks(
          network=network,
          log_prob=log_prob,
          entropy=entropy,
          sample=sample,
          sample_eval=sample_eval,
      )
  
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 4 for calling a global function with example usage: make_mvn_diag_ppo_networks(network) and returns: PPONetworks, tfd, tfd, tfd, tfd