code: def check_batch_sizes(train_results):
      """Check if batch sizes are according to what we expect from config."""
      info = train_results["info"]
      learner_info = info[LEARNER_INFO]
  
      for pid, policy_stats in learner_info.items():
          if pid == "batch_count":
              continue
          # Expect td-errors to be per batch-item.
          configured_b = train_results["config"]["train_batch_size"]
          actual_b = policy_stats["td_error"].shape[0]
          if (configured_b - actual_b) / actual_b > 0.1:
              # Since R2D2 learns on sequences of a fixed length but with variable
              # amount of timesteps that are padded, the batch size is almost never the
              # `train_batch_size`, which is specified in timesteps, but close to it.
              assert 0.8 < (
                  abs(
                      configured_b
                      / (
                          train_results["config"]["model"][
                              "max_seq_len"
                          ]
                          + train_results["config"][
                              "replay_buffer_config"
                          ]["replay_burn_in"]
                      )
                      / actual_b
                  )
                  < 1.2
              )
  
file_path: ray/rllib_contrib/r2d2/tests/test_r2d2.py
imports: import unittest
  from rllib_r2d2.r2d2.r2d2 import R2D2Config
  import ray
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.metrics.learner_info import LEARNER_INFO
  from ray.rllib.utils.test_utils import check_compute_single_action, check_train_results, framework_iterator
label_desc: Use image node_id 1 for calling a global function with example usage: check_batch_sizes(train_results) without return types