code: def test_batch_training(self, name, cls):
      n_arms = 5
      b = cls(n_arms=n_arms)
      total_obs_per_arm = torch.zeros(n_arms)
      total_success_per_arm = torch.zeros(n_arms)
      for _ in range(10):
          n_obs_per_arm = torch.randint(0, 50, size=(n_arms,)).float()
          n_success_per_arm = torch.rand(size=(n_arms,)) * n_obs_per_arm
          total_obs_per_arm += n_obs_per_arm
          total_success_per_arm += n_success_per_arm
  
          b.add_batch_observations(
              n_obs_per_arm,
              n_success_per_arm,
              n_success_per_arm,  # squared rewards are same as rewards
          )
  
          npt.assert_array_equal(
              b.total_n_obs_per_arm.numpy(), total_obs_per_arm.numpy()
          )  # observation counters are correct
          npt.assert_array_equal(
              b.total_sum_reward_per_arm.numpy(),
              total_success_per_arm.numpy(),
          )  # total reward counters are corect
          npt.assert_array_equal(
              b.total_sum_reward_squared_per_arm.numpy(),
              total_success_per_arm.numpy(),
          )  # squared rewards equal to rewards for Bernoulli bandit
  
          self.assertEqual(
              b.total_n_obs_all_arms, total_obs_per_arm.sum().item()
          )  # total observation counter correct
  
          avg_rewards = total_success_per_arm / total_obs_per_arm
          npt.assert_allclose(
              b.get_avg_reward_values().numpy(), avg_rewards.numpy()
          )  # avg rewards computed correctly
  
          scores = b.get_scores()
          forward_scores = b()
  
          # scores shape and type are correct
          self.assertEqual(scores.shape, (n_arms,))
          self.assertIsInstance(scores, torch.Tensor)
          self.assertEqual(forward_scores.shape, (n_arms,))
          self.assertIsInstance(forward_scores, torch.Tensor)
  
          if isinstance(b, BaseUCB):
              npt.assert_array_less(
                  avg_rewards,
                  scores.numpy(),
              )  # UCB scores greater than avg rewards
  
              valid_indices = (
                  b.total_n_obs_per_arm.numpy() >= b.min_num_obs_per_arm
              )
              npt.assert_array_equal(
                  scores[valid_indices], forward_scores[valid_indices]
              )
  
imports: import unittest
  from io import BytesIO
  from itertools import cycle
  from unittest import mock
  import numpy
  import numpy.testing
  import torch
  from parameterized.parameterized import parameterized
  from pytorch_lightning import seed_everything
  from reagent.mab.mab_algorithm import get_arm_indices, place_values_at_indices, randomized_argmax, reindex_multiple_tensors
  from reagent.mab.simulation import BernoilliMAB, compare_bandit_algos, multiple_evaluations_bandit_algo, single_evaluation_bandit_algo
  from reagent.mab.thompson_sampling import BaseThompsonSampling, BernoulliBetaThompson, NormalGammaThompson
  from reagent.mab.ucb import BaseUCB, MetricUCB, UCB1, UCBTuned
label_desc: Use image node_id 2 for calling the TestMAB obj's underlying member method code with example usage: obj.test_batch_training(name, cls) without return types