code: def test_bandit_convergence(self):
      # test whether in a simple bandit environment, the bandit algorithm
      # distribution converge to the optimal distribution empirically
  
      std_threshold = 0.1
      best_arm_prob = 0.5
  
      for config_cls in [BanditLinUCBConfig, BanditLinTSConfig]:
          config = (
              config_cls()
              .debugging(seed=0)
              .environment(
                  env=NonContextualBanditEnv,
                  env_config={"best_arm_prob": best_arm_prob},
              )
          )
          for _ in framework_iterator(
              config,
              frameworks=("tf2", "torch"),
              with_eager_tracing=True,
          ):
              algo = config.build()
              model = algo.get_policy().model
              arm_means, arm_stds = [], []
              for _ in range(50):
                  # TODO the internals of the model is leaking here.
                  # We should revisit this once the RLModule is merged in.
                  samples = [
                      model.arms[i].dist.sample((1000,))
                      for i in range(2)
                  ]
                  arm_means.append(
                      [
                          float(convert_to_numpy(s).mean(0))
                          for s in samples
                      ]
                  )
                  arm_stds.append(
                      [
                          float(convert_to_numpy(s).std(0))
                          for s in samples
                      ]
                  )
                  algo.train()
  
              best_arm = np.argmax(arm_means[-1])
              print(
                  f"best arm: {best_arm}, arm means: {arm_means[-1]}, "
                  f"arm stds: {arm_stds[-1]}"
              )
  
              # the better arm (according to the learned model) should be
              # sufficiently exploited so it should have a low variance at convergence
              self.assertLess(arm_stds[-1][best_arm], std_threshold)
  
              # best arm should also have a good estimate of its actual mean
              # Note that this may not be true for non-optimal arms as they may not
              # have been explored enough
              check(arm_means[-1][best_arm], best_arm_prob, decimals=1)
  
imports: import unittest
  import gymnasium
  import numpy
  from gymnasium.spaces import Box, Discrete
  from rllib_bandit.bandit.bandit import BanditLinTSConfig, BanditLinUCBConfig
  from rllib_bandit.env.bandit_envs_discrete import SimpleContextualBandit
  import ray
  from ray.rllib.env import EnvContext
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.test_utils import check, check_train_results, framework_iterator
label_desc: Use image node_id 5 for calling the TestBandits obj's underlying member method code with example usage: obj.test_bandit_convergence() without return types