code: def policy_network_factory(
      networks: ma_types.MultiAgentNetworks,
      environment_spec: specs.EnvironmentSpec,
      agent_types: Dict[ma_types.AgentID, ma_types.GenericAgent],
      agent_configs: Dict[ma_types.AgentID, ma_types.AgentConfig],
      eval_mode: ma_types.EvalMode,
      init_policy_network_fn: Optional[
          ma_types.InitPolicyNetworkFn
      ] = None,
  ) -> ma_types.MultiAgentPolicyNetworks:
      """Returns default policy networks for all agents.
  
      Args:
        networks: dict mapping agent IDs to their networks.
        environment_spec: environment spec.
        agent_types: dict mapping agent IDs to their type.
        agent_configs: dict mapping agent IDs to their config.
        eval_mode: whether the policy should be initialized in evaluation mode (only
          used if an init_policy_network_fn is not explicitly supplied).
        init_policy_network_fn: optional callable that handles the policy network
          initialization for all sub-agents. If this is not supplied, a default
          policy network initializer is used (if it is supported for the designated
          agent type).
      """
      init_fn = init_policy_network_fn or init_default_policy_network
      policy_networks = {}
      for agent_id, agent_type in agent_types.items():
          single_agent_spec = ma_utils.get_agent_spec(
              environment_spec, agent_id
          )
          policy_networks[agent_id] = init_fn(
              agent_type,
              networks[agent_id],
              single_agent_spec,
              agent_configs[agent_id],
              eval_mode,
          )
      return policy_networks
  
file_path: acme/acme/agents/jax/multiagent/decentralized/factories.py
imports: import enum
  from typing import Any, Dict, Optional
  from acme import specs
  from acme.adders import reverb
  from acme.agents.jax import builders
  from acme.agents.jax import ppo
  from acme.agents.jax import sac
  from acme.agents.jax import td3
  from acme.multiagent import types
  from acme.multiagent import utils
label_desc: Use image node_id 7 for calling a global function with example usage: policy_network_factory(networks, environment_spec, agent_types, agent_configs, eval_mode, init_policy_network_fn) and returns: policy_networks