code: def make_inference_fn(
      ppo_networks: PPONetworks, evaluation: bool = False
  ) -> actor_core_lib.FeedForwardPolicyWithExtra:
      """Returns a function to be used for inference by a PPO actor."""
  
      def inference(
          params: networks_lib.Params,
          key: networks_lib.PRNGKey,
          observations: networks_lib.Observation,
      ):
          dist_params, _ = ppo_networks.network.apply(
              params.model_params, observations
          )
          if evaluation and ppo_networks.sample_eval:
              actions = ppo_networks.sample_eval(dist_params, key)
          else:
              actions = ppo_networks.sample(dist_params, key)
          if evaluation:
              return actions, {}
          log_prob = ppo_networks.log_prob(dist_params, actions)
          extras = {
              "log_prob": log_prob,
              # Add batch dimension.
              "params_num_sgd_steps": params.num_sgd_steps[None, ...],
          }
          return actions, extras
  
      return inference
  
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 1 for calling a global function with example usage: make_inference_fn(ppo_networks, evaluation) and returns: inference, actions,  extras, actions,  dict