code: def make_categorical_ppo_networks(
      network: networks_lib.FeedForwardNetwork,
  ) -> ppo.PPONetworks:
      """Constructs a PPONetworks for Categorical Policy from 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: CategoricalParams, action):
          return tfd.Categorical(logits=params.logits).log_prob(action)
  
      def entropy(params: CategoricalParams, key: networks_lib.PRNGKey):
          del key
          return tfd.Categorical(logits=params.logits).entropy()
  
      def sample(params: CategoricalParams, key: networks_lib.PRNGKey):
          return tfd.Categorical(logits=params.logits).sample(seed=key)
  
      def sample_eval(
          params: CategoricalParams, key: networks_lib.PRNGKey
      ):
          del key
          return tfd.Categorical(logits=params.logits).mode()
  
      return ppo.PPONetworks(
          network=network,
          log_prob=log_prob,
          entropy=entropy,
          sample=sample,
          sample_eval=sample_eval,
      )
  
file_path: acme/examples/multiagent/multigrid/helpers.py
imports: import functools
  from typing import Any, Dict, NamedTuple, Sequence
  from acme import specs
  from acme.agents.jax import ppo
  from acme.agents.jax.multiagent.decentralized import factories
  from acme.jax import networks
  from acme.jax import utils
  from acme.multiagent import types
  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_categorical_ppo_networks(network) and returns: ppo, tfd, tfd, tfd, tfd