code: def test_linucb_discount_factors(self) -> None:
      # change the precision to double
      torch.set_default_dtype(torch.float64)
  
      gamma = 0.8
      scorer = LinearRegressionUCB(self.x_dim, gamma=gamma)
      policy = Policy(scorer=scorer, sampler=GreedyActionSampler())
      trainer = LinUCBTrainer(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.num_arms, self.x_dim)
          ),
          reward=torch.randn((10, 1)),
          action=torch.tensor([[0], [1]], dtype=torch.long).repeat(
              5, 1
          ),
      )
      self.second_batch_2nd_round = CBInput(
          context_arm_features=torch.randn(
              (6, self.num_arms, self.x_dim)
          ),
          reward=torch.randn((6, 1)),
          action=torch.tensor([[0], [1]], dtype=torch.long).repeat(
              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))
      model_output1 = scorer(inp1)
      out1 = model_output1["ucb"]
      # check it won't do redundant coefficent update after 2nd eval
      inp2 = torch.randn((5, self.x_dim))
      model_output2 = scorer(inp2)
      out2 = model_output2["ucb"]
  
      # the feature matrix and model parameter and eval output (computed by hand)
      x1 = add_chosen_arm_features(
          self.batch
      ).features_of_chosen_arm.numpy()
      x2 = add_chosen_arm_features(
          self.batch_2nd_round
      ).features_of_chosen_arm.numpy()
      x3 = add_chosen_arm_features(
          self.second_batch_2nd_round
      ).features_of_chosen_arm.numpy()
      reward1 = self.batch.reward.squeeze().numpy()
      reward2 = self.batch_2nd_round.reward.squeeze().numpy()
      reward3 = self.second_batch_2nd_round.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.avg_A * scorer.sum_weight).numpy(),
          A * scorer.gamma,
          atol=1e-5,
          rtol=1e-5,
      )
      npt.assert_allclose(
          (scorer.avg_b * scorer.sum_weight).numpy(),
          b * scorer.gamma,
          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_avg_A / scorer.sum_weight * scorer.gamma).numpy(),
          inv_A,
          atol=1e-4,
          rtol=1e-4,
      )
  
      # model parameters are the same
      theta = inv_A @ b
      npt.assert_allclose(
          scorer.coefs.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 / scorer.gamma
              )
          return expected_out
  
      expected_out1 = calculated_expected_ucb_scores(inp1)
      npt.assert_allclose(
          out1.numpy(), expected_out1, atol=1e-4, rtol=1e-4
      )
  
      expected_out2 = calculated_expected_ucb_scores(inp2)
      npt.assert_allclose(
          out2.numpy(), expected_out2, atol=1e-4, rtol=1e-4
      )
  
imports: import copy
  import unittest
  import numpy
  import numpy.testing
  import torch
  from libfb.py.collections import AttrDict
  from pytorch_lightning.loggers import TensorBoardLogger
  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.linear_regression import LinearRegressionUCB
  from reagent.training.cb.linucb_trainer import LinUCBTrainer
  from reagent.training.cb.utils import add_chosen_arm_features
  from reagent.training.parameters import LinUCBTrainerParameters
  from torchrec.fb.metrics.meta_config import DefaultTaskInfo, MetricsConfig, QPSDef, RecMetricDef, RecMetricEnum
  from torchrec.fb.metrics.meta_metrics import generate_metric_module, RecMetricModule
label_desc: Use image node_id 7 for calling the TestLinUCB obj's underlying member method code with example usage: obj.test_linucb_discount_factors() and returns: expected_out