code: def test_linucb_discount_factors(self) -> None:
      # change the precision to double
      torch.set_default_dtype(torch.float64)
  
      scorer = DisjointLinearRegressionUCB(
          self.num_arms, self.x_dim, gamma=0.9
      )
      policy = Policy(scorer=scorer, sampler=GreedyActionSampler())
      trainer = DisjointLinUCBTrainer(policy)
  
      # 1st round training
      trainer.training_step(self.batch, 0)
      trainer.on_train_epoch_end()
  
      # 2nd round training
      torch.manual_seed(0)
      self.batch_2nd_round = [
          CBInput(
              context_arm_features=torch.randn((10, self.x_dim)),
              reward=torch.randn((10, 1)),
          ),
          CBInput(
              context_arm_features=torch.randn((3, self.x_dim)),
              reward=torch.randn((3, 1)),
          ),
      ]
      self.second_batch_2nd_round = [
          CBInput(
              context_arm_features=torch.randn((10, self.x_dim)),
              reward=torch.randn((10, 1)),
          ),
          CBInput(
              context_arm_features=torch.randn((3, self.x_dim)),
              reward=torch.randn((3, 1)),
          ),
      ]
      trainer.training_step(self.batch_2nd_round, 0)
      # check if there are several training steps in a round
      # discount factor is only applied one time
      trainer.training_step(self.second_batch_2nd_round, 1)
      trainer.on_train_epoch_end()
  
      # eval dataset
      inp1 = torch.randn((5, self.x_dim))
      out1 = scorer(inp1)
      # check it won't do redundant coefficent update after 2nd eval
      inp2 = torch.randn((5, self.x_dim))
      out2 = scorer(inp2)
  
      # the feature matrix and model parameter and eval output (computed by hand)
      for arm in range(self.num_arms):
          x1 = self.batch[arm].context_arm_features.numpy()
          x2 = self.batch_2nd_round[arm].context_arm_features.numpy()
          x3 = self.second_batch_2nd_round[
              arm
          ].context_arm_features.numpy()
          reward1 = self.batch[arm].reward.squeeze().numpy()
          reward2 = self.batch_2nd_round[arm].reward.squeeze().numpy()
          reward3 = (
              self.second_batch_2nd_round[arm].reward.squeeze().numpy()
          )
  
          # all matrix and vectors are the same
          A = scorer.gamma * x1.T @ x1 + x2.T @ x2 + x3.T @ x3
          b = (
              scorer.gamma * x1.T @ reward1
              + x2.T @ reward2
              + x3.T @ reward3
          )
          npt.assert_allclose(
              scorer.A[arm].numpy(),
              scorer.gamma * A,
              atol=1e-5,
              rtol=1e-5,
          )
          npt.assert_allclose(
              scorer.b[arm].numpy(),
              scorer.gamma * b,
              atol=1e-5,
              rtol=1e-5,
          )
  
          inv_A = np.linalg.inv(
              A + np.identity(self.x_dim) * scorer.l2_reg_lambda
          )
          npt.assert_allclose(
              scorer.inv_A[arm].numpy(), inv_A, atol=1e-4, rtol=1e-4
          )
  
          # model parameters are the same
          theta = inv_A @ b
          npt.assert_allclose(
              scorer.coefs[arm].numpy(), theta, atol=1e-4, rtol=1e-4
          )
  
          # ucb scores are the same
          def calculated_expected_ucb_scores(inp):
              expected_out = np.zeros(inp.size()[0])
              for i in range(inp.size()[0]):
                  x = inp[i].numpy()
                  expected_out[
                      i
                  ] = x @ theta.T + scorer.ucb_alpha * np.sqrt(
                      x @ inv_A @ x.T
                  )
              return expected_out
  
          expected_out1 = calculated_expected_ucb_scores(inp1)
          npt.assert_allclose(
              out1[:, arm].numpy(), expected_out1, atol=1e-4, rtol=1e-4
          )
  
          expected_out2 = calculated_expected_ucb_scores(inp2)
          npt.assert_allclose(
              out2[:, arm].numpy(), expected_out2, atol=1e-4, rtol=1e-4
          )
  
file_path: ReAgent/reagent/test/training/cb/test_disjoint_linucb.py
imports: import copy
  import unittest
  import numpy
  import numpy.testing
  import torch
  from reagent.core.types import CBInput
  from reagent.gym.policies.policy import Policy
  from reagent.gym.policies.samplers.discrete_sampler import GreedyActionSampler
  from reagent.models.disjoint_linucb_predictor import DisjointLinearRegressionUCB
  from reagent.training.cb.disjoint_linucb_trainer import DisjointLinUCBTrainer
  from reagent.training.parameters import DisjointLinUCBTrainerParameters
label_desc: Use image node_id 6 for calling the TestDisjointLinUCB obj's underlying member method code with example usage: obj.test_linucb_discount_factors() and returns: expected_out