docstr: Tests, whether the different model sizes match the ones reported in [1].
docstr_tok: ["Tests",",","whether","the","different","model","sizes","match","the","ones","reported","in","[","1","]","."]
code: def test_dreamerv3_dreamer_model_sizes(self):
      """Tests, whether the different model sizes match the ones reported in [1]."""
  
      # For Atari, these are the exact numbers from the repo ([3]).
      # However, for CartPole + size "S" and "M", the author's original code will not
      # match for the world model count. This is due to the fact that the author uses
      # encoder/decoder nets with 5x1024 nodes (which corresponds to XL) regardless of
      # the `model_size` settings (iff >="S").
      expected_num_params_world_model = {
          "XS_cartpole": 2435076,
          "S_cartpole": 7493380,
          "M_cartpole": 16206084,
          "L_cartpole": 37802244,
          "XL_cartpole": 108353796,
          "XS_atari": 7538979,
          "S_atari": 15687811,
          "M_atari": 32461635,
          "L_atari": 68278275,
          "XL_atari": 181558659,
      }
  
      # All values confirmed against [3] (100% match).
      expected_num_params_actor = {
          # hidden=[1280, 256]
          # hidden_norm=[256], [256]
          # pi (2 actions)=[256, 2], [2]
          "XS_cartpole": 328706,
          "S_cartpole": 1051650,
          "M_cartpole": 2135042,
          "L_cartpole": 4136450,
          "XL_cartpole": 9449474,
          "XS_atari": 329734,
          "S_atari": 1053702,
          "M_atari": 2137606,
          "L_atari": 4139526,
          "XL_atari": 9453574,
      }
  
      # All values confirmed against [3] (100% match).
      expected_num_params_critic = {
          # hidden=[1280, 256]
          # hidden_norm=[256], [256]
          # vf (buckets)=[256, 255], [255]
          "XS_cartpole": 393727,
          "S_cartpole": 1181439,
          "M_cartpole": 2297215,
          "L_cartpole": 4331007,
          "XL_cartpole": 9708799,
          "XS_atari": 393727,
          "S_atari": 1181439,
          "M_atari": 2297215,
          "L_atari": 4331007,
          "XL_atari": 9708799,
      }
  
      config = dreamerv3.DreamerV3Config().training(
          batch_length_T=16,
          horizon_H=5,
          symlog_obs=True,
      )
  
      for _ in framework_iterator(config, frameworks="tf2"):
          # Check all model_sizes described in the paper ([1]) on matching the number
          # of parameters to RLlib's implementation.
          for model_size in ["XS", "S", "M", "L", "XL"]:
              config.model_size = model_size
              config.training(model={"model_size": model_size})
  
              # Atari and CartPole spaces.
              for obs_space, num_actions, env_name in [
                  (
                      gym.spaces.Box(-1.0, 0.0, (4,), np.float32),
                      2,
                      "cartpole",
                  ),
                  (
                      gym.spaces.Box(
                          -1.0, 0.0, (64, 64, 3), np.float32
                      ),
                      6,
                      "atari",
                  ),
              ]:
                  print(
                      f"Testing model_size={model_size} on env-type: {env_name} .."
                  )
                  config.environment(
                      observation_space=obs_space,
                      action_space=gym.spaces.Discrete(num_actions),
                  )
  
                  # Create our RLModule to compute actions with.
                  policy_dict, _ = config.get_multi_agent_setup()
                  module_spec = config.get_marl_module_spec(
                      policy_dict=policy_dict
                  )
                  rl_module = module_spec.build()[DEFAULT_POLICY_ID]
  
                  # Count the generated RLModule's parameters and compare to the
                  # paper's reported numbers ([1] and [3]).
                  num_params_world_model = sum(
                      np.prod(v.shape.as_list())
                      for v in rl_module.world_model.trainable_variables
                  )
                  self.assertEqual(
                      num_params_world_model,
                      expected_num_params_world_model[
                          f"{model_size}_{env_name}"
                      ],
                  )
                  num_params_actor = sum(
                      np.prod(v.shape.as_list())
                      for v in rl_module.actor.trainable_variables
                  )
                  self.assertEqual(
                      num_params_actor,
                      expected_num_params_actor[
                          f"{model_size}_{env_name}"
                      ],
                  )
                  num_params_critic = sum(
                      np.prod(v.shape.as_list())
                      for v in rl_module.critic.trainable_variables
                  )
                  self.assertEqual(
                      num_params_critic,
                      expected_num_params_critic[
                          f"{model_size}_{env_name}"
                      ],
                  )
                  print("\tok")
  
file_path: ray/rllib/algorithms/dreamerv3/tests/test_dreamerv3.py
imports: import unittest
  import gymnasium
  import numpy
  import ray
  from ray.rllib.algorithms.dreamerv3 import dreamerv3
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
  from ray.rllib.utils.numpy import one_hot
  from ray.rllib.utils.test_utils import framework_iterator
  from ray import tune
label_desc: Use image node_id 4 for calling the TestDreamerV3 obj's underlying member method code with example usage: obj.test_dreamerv3_dreamer_model_sizes() without return types