code: def make_ppo_networks(
      network: networks_lib.FeedForwardNetwork,
  ) -> PPONetworks:
      """Constructs a PPONetworks instance from the given FeedForwardNetwork.
  
      This method assumes that the network returns a tfd.Distribution. Sometimes it
      may be preferable to have networks that do not return tfd.Distribution
      objects, for example, due to tfd.Distribution not playing nice with jax.vmap.
      Please refer to the make_continuous_networks() for an example where the
      network does not return a tfd.Distribution object.
  
      Args:
        network: a transformed Haiku network that takes in observations and returns
          the action distribution and value.
  
      Returns:
        A PPONetworks instance with pure functions wrapping the input network.
      """
      return PPONetworks(
          network=network,
          log_prob=lambda distribution, action: distribution.log_prob(
              action
          ),
          entropy=lambda distribution, key=None: distribution.entropy(),
          sample=lambda distribution, key: distribution.sample(
              seed=key
          ),
          sample_eval=lambda distribution, key: distribution.mode(),
      )
  
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 3 for calling a global function with example usage: make_ppo_networks(network) and returns: PPONetworks