docstr: Tests if model output spec checking works correctly.
  
  This test is centered around the `always_check_shapes` flag of the
  ModelConfig class. If this flag is set to True, the model will always
  check if the outputs conform to the specs. If this flag is set to False,
  the model will never check the outputs.
docstr_tok: ["Tests","if","model","output","spec","checking","works","correctly",".","This","test","is","centered","around","the","`","always_check_shapes","`","flag","of","the","ModelConfig","class",".","If","this","flag","is","set","to","True",",","the","model","will","always","check","if","the","outputs","conform","to","the","specs",".","If","this","flag","is","set","to","False",",","the","model","will","never","check","the","outputs","."]
code: def test_model_output_spec_checking(self):
      """Tests if model output spec checking works correctly.
  
      This test is centered around the `always_check_shapes` flag of the
      ModelConfig class. If this flag is set to True, the model will always
      check if the outputs conform to the specs. If this flag is set to False,
      the model will never check the outputs.
      """
  
      for fw in ["torch", "tf2"]:
  
          class BadModel:
              """Simple model that produces bad outputs."""
  
              def get_output_specs(self):
                  return SpecDict(
                      {
                          "out": TensorSpec("b, h", h=1),
                      }
                  )
  
              def get_input_specs(self):
                  return SpecDict(
                      {
                          "in": TensorSpec("b, h", h=1),
                      }
                  )
  
          if fw == "tf2":
  
              class TestModel(BadModel, TfModel):
                  def _forward(self, input_dict):
                      return {"out": torch.Tensor([[1, 2]])}
  
          else:
  
              class TestModel(BadModel, TfModel):
                  def _forward(self, input_dict):
                      return {"out": tf.constant([[1, 2]])}
  
          @dataclass
          class CatModelConfig(ModelConfig):
              def build(self, framework: str):
                  # Since we define the correct model above anyway, we don't need
                  # to distinguish between frameworks here.
                  return TestModel(self)
  
          # 1) Check if model behaves correctly with always_check_shapes=True first.
          # We expect model to raise an error if the output shapes are not correct.
          # This is the behaviour we use for debugging with model specs.
  
          config = CatModelConfig(always_check_shapes=True)
  
          model = config.build(framework="spam")
  
          # We want to raise an output spec validation error here since the output
          # has the wrong shape
          with self.assertRaisesRegex(
              SpecCheckingError, "output spec validation failed"
          ):
              model({"in": torch.Tensor([[1]])})
  
          # 2) Check if model behaves correctly with always_check_shapes=False.
          # We don't expect model to raise an error.
          # This is the more performant default behaviour
  
          config = CatModelConfig(always_check_shapes=False)
  
          model = config.build(framework="spam")
  
          model({"in_1": [[1]]})
  
imports: import unittest
  from dataclasses import dataclass
  import gymnasium
  from ray.rllib.core.models.configs import ModelConfig
  from ray.rllib.core.models.specs.checker import SpecCheckingError
  from ray.rllib.core.models.specs.specs_base import TensorSpec
  from ray.rllib.core.models.specs.specs_dict import SpecDict
  from ray.rllib.core.models.tf.base import TfModel
  from ray.rllib.core.models.torch.base import TorchModel
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
  from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
  from ray.rllib.core.rl_module.torch.torch_compile_config import TorchCompileConfig
  from ray.rllib.utils.torch_utils import _dynamo_is_available
label_desc: Use image node_id 2 for calling the TestModelBase obj's underlying member method code with example usage: obj.test_model_output_spec_checking() and returns: SpecDict, SpecDict, TestModel, dict, dict