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: Implementation of the BaseModelComposer interface, returns the value networks.
docstr_tok: ["Implementation","of","the","BaseModelComposer","interface",",","returns","the","value","networks","."]
code: def critic(
      self,
  ) -> Optional[Union[TorchStateCritic, TorchStateActionCritic]]:
      """Implementation of the BaseModelComposer interface, returns the value networks."""
  
      if self._critic_type is None:
          return None
  
      elif issubclass(self._critic_type, SharedStateCriticComposer):
          assert not any(
              self.model_builder.use_shared_embedding.values()
          ), f"Embedding can not be shared when using a shared state critic"
          observation_space = flat_structured_space(
              self.observation_spaces_dict
          )
          critics = {0: self.template_value_net(observation_space)}
          return TorchSharedStateCritic(
              networks=critics,
              obs_spaces_dict=self.observation_spaces_dict,
              device="cpu",
              stack_observations=False,
          )
  
      elif issubclass(self._critic_type, StepStateCriticComposer):
          critics = dict()
          for (
              sub_step_key,
              sub_step_space,
          ) in self.observation_spaces_dict.items():
              critics[sub_step_key] = self.template_value_net(
                  observation_space=sub_step_space,
                  perception_net=self._shared_embedding_nets[
                      sub_step_key
                  ],
                  shared_embedding_keys=self.model_builder.shared_embedding_keys[
                      sub_step_key
                  ],
              )
          return TorchStepStateCritic(
              networks=critics,
              obs_spaces_dict=self.observation_spaces_dict,
              device="cpu",
          )
  
      elif issubclass(self._critic_type, DeltaStateCriticComposer):
          critics = dict()
          for idx, (sub_step_key, sub_step_space) in enumerate(
              self.observation_spaces_dict.items()
          ):
              if idx > 0:
                  sub_step_space = spaces.Dict(
                      {
                          **sub_step_space.spaces,
                          **DeltaStateCriticComposer.prev_value_space.spaces,
                      }
                  )
              critics[sub_step_key] = self.template_value_net(
                  observation_space=sub_step_space,
                  perception_net=self._shared_embedding_nets[
                      sub_step_key
                  ],
                  shared_embedding_keys=self.model_builder.shared_embedding_keys[
                      sub_step_key
                  ],
              )
  
          return TorchDeltaStateCritic(
              networks=critics,
              obs_spaces_dict=self.observation_spaces_dict,
              device="cpu",
          )
      elif issubclass(
          self._critic_type, SharedStateActionCriticComposer
      ):
          assert not any(
              self.model_builder.use_shared_embedding.values()
          ), f"Embedding can not be shared when using a shared state action critic"
  
          observation_space = flat_structured_space(
              self.observation_spaces_dict
          )
          action_space = flat_structured_space(self.action_spaces_dict)
          only_discrete_spaces = self._only_discrete_spaces()
          critics = {
              0: self.template_q_value_net(
                  observation_space,
                  action_space,
                  all(only_discrete_spaces.values()),
              )
          }
          return TorchSharedStateActionCritic(
              networks=critics,
              num_policies=len(self.action_spaces_dict),
              device="cpu",
              only_discrete_spaces=only_discrete_spaces,
              action_spaces_dict={0: action_space},
          )
  
      elif issubclass(self._critic_type, StepStateActionCriticComposer):
          assert not any(
              self.model_builder.use_shared_embedding.values()
          ), f"Embedding can not be shared when using a step state action critic"
          critics = dict()
          only_discrete_spaces = self._only_discrete_spaces()
          for (
              sub_step_key,
              sub_step_space,
          ) in self.observation_spaces_dict.items():
              critics[sub_step_key] = self.template_q_value_net(
                  observation_space=sub_step_space,
                  action_space=self.action_spaces_dict[sub_step_key],
                  only_discrete_spaces=only_discrete_spaces[
                      sub_step_key
                  ],
              )
          return TorchStepStateActionCritic(
              networks=critics,
              num_policies=len(self.action_spaces_dict),
              device="cpu",
              only_discrete_spaces=only_discrete_spaces,
              action_spaces_dict=self.action_spaces_dict,
          )
  
      else:
          raise ValueError(
              f"Critics type {self._critic_type} not supported by the template model composer!"
          )
  
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 9 for calling the TemplateModelComposer obj's underlying member method code with example usage: obj.critic() and returns: None, TorchSharedStateCritic, TorchStepStateCritic, TorchDeltaStateCritic, TorchSharedStateActionCritic, TorchStepStateActionCritic