class_docstr: Provides a mapping of spaces and action heads to the respective probability distributions to be used.
  
  This ensures full flexibility for specifying different distributions to the same gym action space type.
  (e.g. One gym.spaces.Box space could be modeled with a Beta another one with a DiagonalGaussian distribution.)
  It allows to add and register arbitrary custom distributions.
  
  :param action_space: The dictionary action space.
  :param distribution_mapper_config: A Distribution mapper configuration (for details see the docs).
class_docstr_tok: ["Provides","a","mapping","of","spaces","and","action","heads","to","the","respective","probability","distributions","to","be","used",".","This","ensures","full","flexibility","for","specifying","different","distributions","to","the","same","gym","action","space","type",".","(","e.g",".","One","gym.spaces.Box","space","could","be","modeled","with","a","Beta","another","one","with","a","DiagonalGaussian","distribution",".",")","It","allows","to","add","and","register","arbitrary","custom","distributions",".",":","param","action_space",":","The","dictionary","action","space",".",":","param","distribution_mapper_config",":","A","Distribution","mapper","configuration","(","for","details","see","the","docs",")","."]
code: def __init__(
      self,
      action_space: spaces.Dict,
      distribution_mapper_config: CollectionOfConfigType,
  ):
      self.action_space = action_space
  
      # mapping of action heads to distributions and configs
      self._action_head_to_distribution: Dict[
          str, Tuple[type(TorchProbabilityDistribution), Dict[str, Any]]
      ] = dict()
  
      # first: apply default config to action heads
      for action_head, sub_action_space in action_space.spaces.items():
          space_type = type(sub_action_space)
          dist_type: TorchProbabilityDistribution = (
              self.default_mapping[space_type]
          )
          self._action_head_to_distribution[action_head] = (
              dist_type,
              {},
          )
  
      # second: parse custom mappings
      for entry_dict in distribution_mapper_config:
          assert "distribution" in entry_dict
          assert (
              "action_space" in entry_dict
              and "action_head" not in entry_dict
          ) or (
              "action_space" not in entry_dict
              and "action_head" in entry_dict
          )
  
          # get the distribution type
          distribution_type = Factory(
              TorchProbabilityDistribution
          ).type_from_name(entry_dict["distribution"])
  
          # get additional distribution arguments
          args = entry_dict["args"] if "args" in entry_dict else {}
  
          if "action_head" in entry_dict:
              self._action_head_to_distribution[
                  entry_dict["action_head"]
              ] = (distribution_type, args)
  
          elif "action_space" in entry_dict:
              sub_action_space = Factory(spaces.Space).type_from_name(
                  entry_dict["action_space"]
              )
  
              for action_head in self.action_space.spaces:
                  if isinstance(
                      self.action_space[action_head], sub_action_space
                  ):
                      self._action_head_to_distribution[action_head] = (
                          distribution_type,
                          args,
                      )
  
file_path: maze/maze/distributions/distribution_mapper.py
imports: from typing import Dict, Tuple, Any, Sequence
  import torch
  from gym import spaces
  from maze.core.utils.factory import Factory, CollectionOfConfigType
  from maze.distributions.bernoulli import BernoulliProbabilityDistribution
  from maze.distributions.categorical import CategoricalProbabilityDistribution
  from maze.distributions.dict import DictProbabilityDistribution
  from maze.distributions.gaussian import DiagonalGaussianProbabilityDistribution
  from maze.distributions.multi_categorical import MultiCategoricalProbabilityDistribution
  from maze.distributions.torch_dist import TorchProbabilityDistribution
label_desc: Use image node_id 1 to create a new DistributionMapper object  with example: obj = DistributionMapper(action_space, distribution_mapper_config)