class_docstr: Composes template models from configs.
  
  :param action_spaces_dict: Dict of sub-step id to action space.
  :param observation_spaces_dict: Dict of sub-step id to observation space.
  :param distribution_mapper_config: Distribution mapper configuration.
  :param model_builder: The model builder (template) to create the model from.
  :param policy: Specifies the policy type as a configType.
      E.g. {'type': maze.perception.models.policies.ProbabilisticPolicyComposer} specifies a probabilistic policy.
  :param critic: Specifies the critic type as a configType.
      E.g. {'type': maze.perception.models.critics.StateCriticComposer} specifies the single step state critic.
class_docstr_tok: ["Composes","template","models","from","configs",".",":","param","action_spaces_dict",":","Dict","of","sub-step","id","to","action","space",".",":","param","observation_spaces_dict",":","Dict","of","sub-step","id","to","observation","space",".",":","param","distribution_mapper_config",":","Distribution","mapper","configuration",".",":","param","model_builder",":","The","model","builder","(","template",")","to","create","the","model","from",".",":","param","policy",":","Specifies","the","policy","type","as","a","configType",".","E.g",".","{","'type","'",":","maze.perception.models.policies.ProbabilisticPolicyComposer","}","specifies","a","probabilistic","policy",".",":","param","critic",":","Specifies","the","critic","type","as","a","configType",".","E.g",".","{","'type","'",":","maze.perception.models.critics.StateCriticComposer","}","specifies","the","single","step","state","critic","."]
docstr: Compiles a template policy network.
  
  :param observation_space: The input observations for the perception network.
  :param action_space: The action space that defines the network action heads.
  :param shared_embedding_keys: The list of embedding keys for this substep's model.
  :return: A policy network (actor) InferenceBlock, as well as the embedding net InferenceBlock if shared keys
      have been specified.
docstr_tok: ["Compiles","a","template","policy","network",".",":","param","observation_space",":","The","input","observations","for","the","perception","network",".",":","param","action_space",":","The","action","space","that","defines","the","network","action","heads",".",":","param","shared_embedding_keys",":","The","list","of","embedding","keys","for","this","substep","'s","model",".",":","return",":","A","policy","network","(","actor",")","InferenceBlock",",","as","well","as","the","embedding","net","InferenceBlock","if","shared","keys","have","been","specified","."]
code: def template_policy_net(
      self,
      observation_space: spaces.Dict,
      action_space: spaces.Dict,
      shared_embedding_keys: List[str],
  ) -> Tuple[InferenceBlock, InferenceBlock]:
      """Compiles a template policy network.
  
      :param observation_space: The input observations for the perception network.
      :param action_space: The action space that defines the network action heads.
      :param shared_embedding_keys: The list of embedding keys for this substep's model.
      :return: A policy network (actor) InferenceBlock, as well as the embedding net InferenceBlock if shared keys
          have been specified.
      """
  
      # build perception net
      embedding_net = self.template_perception_net(observation_space)
  
      # build action head
      perception_dict = embedding_net.perception_dict
      action_heads = []
      for action_head, action_space in action_space.spaces.items():
          # initialize action head
          action_net = LinearOutputBlock(
              in_keys="latent",
              out_keys=action_head,
              in_shapes=perception_dict["latent"].out_shapes(),
              output_units=self._distribution_mapper.required_logits_shape(
                  action_head
              )[
                  0
              ],
          )
  
          module_init = make_module_init_normc(std=0.01)
          action_net.apply(module_init)
  
          # extent perception dictionary
          perception_dict[action_head] = action_net
          action_heads.append(action_head)
  
      # compile inference model
      shared_embedding_keys_remove_input = list(
          filter(
              lambda x: x not in embedding_net.in_keys,
              shared_embedding_keys,
          )
      )
      net = InferenceBlock(
          in_keys=embedding_net.in_keys,
          out_keys=action_heads + shared_embedding_keys_remove_input,
          in_shapes=embedding_net.in_shapes,
          perception_blocks=perception_dict,
      )
  
      if len(shared_embedding_keys_remove_input) == 0:
          embedding_net = None
  
      return net, embedding_net
  
imports: import functools
  from typing import Optional, Dict, Union, Type, Tuple, List
  from gym import spaces
  from maze.core.agent.torch_policy import TorchPolicy
  from maze.core.agent.torch_state_action_critic import TorchSharedStateActionCritic, TorchStateActionCritic, TorchStepStateActionCritic
  from maze.core.agent.torch_state_critic import TorchStateCritic, TorchSharedStateCritic, TorchStepStateCritic, TorchDeltaStateCritic
  from maze.core.annotations import override
  from maze.core.env.structured_env import StepKeyType
  from maze.core.utils.factory import ConfigType, Factory
  from maze.core.utils.structured_env_utils import flat_structured_space
  from maze.core.wrappers.observation_preprocessing.preprocessors.one_hot import OneHotPreProcessor
  from maze.perception.blocks.inference import InferenceBlock
  from maze.perception.blocks.output.linear import LinearOutputBlock
  from maze.perception.builders.base import BaseModelBuilder
  from maze.perception.models.critics import SharedStateCriticComposer, DeltaStateCriticComposer, StepStateCriticComposer
  from maze.perception.models.critics.critic_composer_interface import CriticComposerInterface
  from maze.perception.models.critics.shared_state_action_critics_composer import SharedStateActionCriticComposer
  from maze.perception.models.critics.step_state_action_critic_composer import StepStateActionCriticComposer
  from maze.perception.models.model_composer import BaseModelComposer
  from maze.perception.models.policies.base_policy_composer import BasePolicyComposer
  from maze.perception.models.policies.probabilistic_policy_composer import ProbabilisticPolicyComposer
  from maze.perception.weight_init import make_module_init_normc
  from maze.utils.bcolors import BColors
label_desc: Use image node_id 4 for calling the TemplateModelComposer obj's underlying member method code with example usage: obj.template_policy_net(observation_space, action_space, shared_embedding_keys) and returns: net,  embedding_net