class_docstr: Registry of models, preprocessors, and action distributions for envs.
  
  .. testcode::
      :skipif: True
  
      prep = ModelCatalog.get_preprocessor(env)
      observation = prep.transform(raw_observation)
  
      dist_class, dist_dim = ModelCatalog.get_action_dist(
          env.action_space, {})
      model = ModelCatalog.get_model_v2(
          obs_space, action_space, num_outputs, options)
      dist = dist_class(model.outputs, model)
      action = dist.sample()
class_docstr_tok: ["Registry","of","models",",","preprocessors",",","and","action","distributions","for","envs",".","..","testcode",":",":",":","skipif",":","True","prep","=","ModelCatalog.get_preprocessor","(","env",")","observation","=","prep.transform","(","raw_observation",")","dist_class",",","dist_dim","=","ModelCatalog.get_action_dist","(","env.action_space",",","{","}",")","model","=","ModelCatalog.get_model_v2","(","obs_space",",","action_space",",","num_outputs",",","options",")","dist","=","dist_class","(","model.outputs",",","model",")","action","=","dist.sample","(",")"]
docstr: Returns a suitable model compatible with given spaces and output.
  
  Args:
      obs_space: Observation space of the target gym env. This
          may have an `original_space` attribute that specifies how to
          unflatten the tensor into a ragged tensor.
      action_space: Action space of the target gym env.
      num_outputs: The size of the output vector of the model.
      model_config: The "model" sub-config dict
          within the Algorithm's config dict.
      framework: One of "tf2", "tf", "torch", or "jax".
      name: Name (scope) for the model.
      model_interface: Interface required for the model
      default_model: Override the default class for the model. This
          only has an effect when not using a custom model
      model_kwargs: Args to pass to the ModelV2 constructor
  
  Returns:
      model (ModelV2): Model to use for the policy.
docstr_tok: ["Returns","a","suitable","model","compatible","with","given","spaces","and","output",".","Args",":","obs_space",":","Observation","space","of","the","target","gym","env",".","This","may","have","an","`","original_space","`","attribute","that","specifies","how","to","unflatten","the","tensor","into","a","ragged","tensor",".","action_space",":","Action","space","of","the","target","gym","env",".","num_outputs",":","The","size","of","the","output","vector","of","the","model",".","model_config",":","The","``","model","''","sub-config","dict","within","the","Algorithm","'s","config","dict",".","framework",":","One","of","``","tf2","''",",","``","tf","''",",","``","torch","''",",","or","``","jax","''",".","name",":","Name","(","scope",")","for","the","model",".","model_interface",":","Interface","required","for","the","model","default_model",":","Override","the","default","class","for","the","model",".","This","only","has","an","effect","when","not","using","a","custom","model","model_kwargs",":","Args","to","pass","to","the","ModelV2","constructor","Returns",":","model","(","ModelV2",")",":","Model","to","use","for","the","policy","."]
code: def get_model_v2(
      obs_space: gym.Space,
      action_space: gym.Space,
      num_outputs: int,
      model_config: ModelConfigDict,
      framework: str = "tf",
      name: str = "default_model",
      model_interface: type = None,
      default_model: type = None,
      **model_kwargs,
  ) -> ModelV2:
      """Returns a suitable model compatible with given spaces and output.
  
      Args:
          obs_space: Observation space of the target gym env. This
              may have an `original_space` attribute that specifies how to
              unflatten the tensor into a ragged tensor.
          action_space: Action space of the target gym env.
          num_outputs: The size of the output vector of the model.
          model_config: The "model" sub-config dict
              within the Algorithm's config dict.
          framework: One of "tf2", "tf", "torch", or "jax".
          name: Name (scope) for the model.
          model_interface: Interface required for the model
          default_model: Override the default class for the model. This
              only has an effect when not using a custom model
          model_kwargs: Args to pass to the ModelV2 constructor
  
      Returns:
          model (ModelV2): Model to use for the policy.
      """
  
      # Validate the given config dict.
      ModelCatalog._validate_config(
          config=model_config,
          action_space=action_space,
          framework=framework,
      )
  
      if model_config.get("custom_model"):
          # Allow model kwargs to be overridden / augmented by
          # custom_model_config.
          customized_model_kwargs = dict(
              model_kwargs,
              **model_config.get("custom_model_config", {}),
          )
  
          if isinstance(model_config["custom_model"], type):
              model_cls = model_config["custom_model"]
          elif (
              isinstance(model_config["custom_model"], str)
              and "." in model_config["custom_model"]
          ):
              return from_config(
                  cls=model_config["custom_model"],
                  obs_space=obs_space,
                  action_space=action_space,
                  num_outputs=num_outputs,
                  model_config=customized_model_kwargs,
                  name=name,
              )
          else:
              model_cls = _global_registry.get(
                  RLLIB_MODEL, model_config["custom_model"]
              )
  
          # Only allow ModelV2 or native keras Models.
          if not issubclass(model_cls, ModelV2):
              if framework not in ["tf", "tf2"] or not issubclass(
                  model_cls, tf.keras.Model
              ):
                  raise ValueError(
                      "`model_cls` must be a ModelV2 sub-class, but is"
                      " {}!".format(model_cls)
                  )
  
          logger.info(
              "Wrapping {} as {}".format(model_cls, model_interface)
          )
          model_cls = ModelCatalog._wrap_if_needed(
              model_cls, model_interface
          )
  
          if framework in ["tf2", "tf"]:
              # Try wrapping custom model with LSTM/attention, if required.
              if model_config.get("use_lstm") or model_config.get(
                  "use_attention"
              ):
                  from ray.rllib.models.tf.attention_net import (
                      AttentionWrapper,
                  )
                  from ray.rllib.models.tf.recurrent_net import (
                      LSTMWrapper,
                  )
  
                  wrapped_cls = model_cls
                  forward = wrapped_cls.forward
                  model_cls = ModelCatalog._wrap_if_needed(
                      wrapped_cls,
                      LSTMWrapper
                      if model_config.get("use_lstm")
                      else AttentionWrapper,
                  )
                  model_cls._wrapped_forward = forward
  
              # Obsolete: Track and warn if vars were created but not
              # registered. Only still do this, if users do register their
              # variables. If not (which they shouldn't), don't check here.
              created = set()
  
              def track_var_creation(next_creator, **kw):
                  v = next_creator(**kw)
                  created.add(v.ref())
                  return v
  
              with tf.variable_creator_scope(track_var_creation):
                  if issubclass(model_cls, tf.keras.Model):
                      instance = model_cls(
                          input_space=obs_space,
                          action_space=action_space,
                          num_outputs=num_outputs,
                          name=name,
                          **customized_model_kwargs,
                      )
                  else:
                      # Try calling with kwargs first (custom ModelV2 should
                      # accept these as kwargs, not get them from
                      # config["custom_model_config"] anymore).
                      try:
                          instance = model_cls(
                              obs_space,
                              action_space,
                              num_outputs,
                              model_config,
                              name,
                              **customized_model_kwargs,
                          )
                      except TypeError as e:
                          # Keyword error: Try old way w/o kwargs.
                          if (
                              "__init__() got an unexpected "
                              in e.args[0]
                          ):
                              instance = model_cls(
                                  obs_space,
                                  action_space,
                                  num_outputs,
                                  model_config,
                                  name,
                                  **model_kwargs,
                              )
                              logger.warning(
                                  "Custom ModelV2 should accept all custom "
                                  "options as **kwargs, instead of expecting"
                                  " them in config['custom_model_config']!"
                              )
                          # Other error -> re-raise.
                          else:
                              raise e
  
              # User still registered TFModelV2's variables: Check, whether
              # ok.
              registered = []
              if not isinstance(instance, tf.keras.Model):
                  registered = set(instance.var_list)
              if len(registered) > 0:
                  not_registered = set()
                  for var in created:
                      if var not in registered:
                          not_registered.add(var)
                  if not_registered:
                      raise ValueError(
                          "It looks like you are still using "
                          "`{}.register_variables()` to register your "
                          "model's weights. This is no longer required, but "
                          "if you are still calling this method at least "
                          "once, you must make sure to register all created "
                          "variables properly. The missing variables are {},"
                          " and you only registered {}. "
                          "Did you forget to call `register_variables()` on "
                          "some of the variables in question?".format(
                              instance, not_registered, registered
                          )
                      )
          elif framework == "torch":
              # Try wrapping custom model with LSTM/attention, if required.
              if model_config.get("use_lstm") or model_config.get(
                  "use_attention"
              ):
                  from ray.rllib.models.torch.attention_net import (
                      AttentionWrapper,
                  )
                  from ray.rllib.models.torch.recurrent_net import (
                      LSTMWrapper,
                  )
  
                  wrapped_cls = model_cls
                  forward = wrapped_cls.forward
                  model_cls = ModelCatalog._wrap_if_needed(
                      wrapped_cls,
                      LSTMWrapper
                      if model_config.get("use_lstm")
                      else AttentionWrapper,
                  )
                  model_cls._wrapped_forward = forward
  
              # PyTorch automatically tracks nn.Modules inside the parent
              # nn.Module's constructor.
              # Try calling with kwargs first (custom ModelV2 should
              # accept these as kwargs, not get them from
              # config["custom_model_config"] anymore).
              try:
                  instance = model_cls(
                      obs_space,
                      action_space,
                      num_outputs,
                      model_config,
                      name,
                      **customized_model_kwargs,
                  )
              except TypeError as e:
                  # Keyword error: Try old way w/o kwargs.
                  if "__init__() got an unexpected " in e.args[0]:
                      instance = model_cls(
                          obs_space,
                          action_space,
                          num_outputs,
                          model_config,
                          name,
                          **model_kwargs,
                      )
                      logger.warning(
                          "Custom ModelV2 should accept all custom "
                          "options as **kwargs, instead of expecting"
                          " them in config['custom_model_config']!"
                      )
                  # Other error -> re-raise.
                  else:
                      raise e
          else:
              raise NotImplementedError(
                  "`framework` must be 'tf2|tf|torch', but is "
                  "{}!".format(framework)
              )
  
          return instance
  
      # Find a default TFModelV2 and wrap with model_interface.
      if framework in ["tf", "tf2"]:
          v2_class = None
          # Try to get a default v2 model.
          if not model_config.get("custom_model"):
              v2_class = (
                  default_model
                  or ModelCatalog._get_v2_model_class(
                      obs_space, model_config, framework=framework
                  )
              )
  
          if not v2_class:
              raise ValueError("ModelV2 class could not be determined!")
  
          if model_config.get("use_lstm") or model_config.get(
              "use_attention"
          ):
              from ray.rllib.models.tf.attention_net import (
                  AttentionWrapper,
              )
              from ray.rllib.models.tf.recurrent_net import (
                  LSTMWrapper,
              )
  
              wrapped_cls = v2_class
              if model_config.get("use_lstm"):
                  v2_class = ModelCatalog._wrap_if_needed(
                      wrapped_cls, LSTMWrapper
                  )
                  v2_class._wrapped_forward = wrapped_cls.forward
              else:
                  v2_class = ModelCatalog._wrap_if_needed(
                      wrapped_cls, AttentionWrapper
                  )
                  v2_class._wrapped_forward = wrapped_cls.forward
  
          # Wrap in the requested interface.
          wrapper = ModelCatalog._wrap_if_needed(
              v2_class, model_interface
          )
  
          if issubclass(wrapper, tf.keras.Model):
              model = wrapper(
                  input_space=obs_space,
                  action_space=action_space,
                  num_outputs=num_outputs,
                  name=name,
                  **dict(model_kwargs, **model_config),
              )
              return model
  
          return wrapper(
              obs_space,
              action_space,
              num_outputs,
              model_config,
              name,
              **model_kwargs,
          )
  
      # Find a default TorchModelV2 and wrap with model_interface.
      elif framework == "torch":
          # Try to get a default v2 model.
          if not model_config.get("custom_model"):
              v2_class = (
                  default_model
                  or ModelCatalog._get_v2_model_class(
                      obs_space, model_config, framework=framework
                  )
              )
  
          if not v2_class:
              raise ValueError("ModelV2 class could not be determined!")
  
          if model_config.get("use_lstm") or model_config.get(
              "use_attention"
          ):
              from ray.rllib.models.torch.attention_net import (
                  AttentionWrapper,
              )
              from ray.rllib.models.torch.recurrent_net import (
                  LSTMWrapper,
              )
  
              wrapped_cls = v2_class
              forward = wrapped_cls.forward
              if model_config.get("use_lstm"):
                  v2_class = ModelCatalog._wrap_if_needed(
                      wrapped_cls, LSTMWrapper
                  )
              else:
                  v2_class = ModelCatalog._wrap_if_needed(
                      wrapped_cls, AttentionWrapper
                  )
  
              v2_class._wrapped_forward = forward
  
          # Wrap in the requested interface.
          wrapper = ModelCatalog._wrap_if_needed(
              v2_class, model_interface
          )
          return wrapper(
              obs_space,
              action_space,
              num_outputs,
              model_config,
              name,
              **model_kwargs,
          )
  
      # Find a default JAXModelV2 and wrap with model_interface.
      elif framework == "jax":
          v2_class = default_model or ModelCatalog._get_v2_model_class(
              obs_space, model_config, framework=framework
          )
          # Wrap in the requested interface.
          wrapper = ModelCatalog._wrap_if_needed(
              v2_class, model_interface
          )
          return wrapper(
              obs_space,
              action_space,
              num_outputs,
              model_config,
              name,
              **model_kwargs,
          )
      else:
          raise NotImplementedError(
              "`framework` must be 'tf2|tf|torch', but is "
              "{}!".format(framework)
          )
  
imports: from functools import partial
  import gymnasium
  from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
  import logging
  import numpy
  import tree
  from typing import List, Optional, Type, Union
  from ray.tune.registry import RLLIB_MODEL, RLLIB_ACTION_DIST, _global_registry
  from ray.rllib.models.action_dist import ActionDistribution
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.preprocessors import get_preprocessor, Preprocessor
  from ray.rllib.models.tf.tf_action_dist import Categorical, Deterministic, DiagGaussian, Dirichlet, MultiActionDistribution, MultiCategorical
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical, TorchDeterministic, TorchDirichlet, TorchDiagGaussian, TorchMultiActionDistribution, TorchMultiCategorical
  from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI
  from ray.rllib.utils.deprecation import DEPRECATED_VALUE, deprecation_warning
  from ray.rllib.utils.error import UnsupportedSpaceException
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.from_config import from_config
  from ray.rllib.utils.spaces.simplex import Simplex
  from ray.rllib.utils.spaces.space_utils import flatten_space
  from ray.rllib.utils.typing import ModelConfigDict, TensorType
label_desc: Use image node_id 4 for calling the ModelCatalog obj's underlying member method code with example usage: obj.get_model_v2(obs_space, action_space, num_outputs, model_config, framework, name, model_interface, default_model) and returns: instance, wrapper, model, wrapper, from_config, v, wrapper