code: def test_compute_discount_tensor(self):
      time_diff = 4
      steps = 3
      inp = DiscreteDqnInput(
          state=FeatureData(
              float_features=torch.rand(self.batch_size, self.state_dim)
          ),
          next_state=FeatureData(
              float_features=torch.rand(self.batch_size, self.state_dim)
          ),
          reward=torch.ones(self.batch_size, 1),
          time_diff=torch.ones(self.batch_size, 1) * time_diff,
          step=torch.ones(self.batch_size, 1) * steps,
          not_terminal=torch.ones(self.batch_size, 1),
          action=torch.tensor([[0, 1], [1, 0], [0, 1]]),
          next_action=torch.tensor([[1, 0], [0, 1], [1, 0]]),
          possible_actions_mask=torch.ones(
              self.batch_size, self.action_dim
          ),
          possible_next_actions_mask=torch.ones(
              self.batch_size, self.action_dim
          ),
          extras=ExtraData(),
      )
  
      # vanilla
      trainer = self._construct_trainer()
      discount_tensor = trainer.compute_discount_tensor(
          batch=inp, boosted_rewards=inp.reward
      )
      self.assertEqual(discount_tensor.shape[0], self.batch_size)
      self.assertEqual(discount_tensor.shape[1], 1)
      self.assertTrue(
          torch.isclose(
              discount_tensor, torch.tensor(trainer.gamma)
          ).all()
      )
  
      # seq_num
      param_copy = DQNTrainerParameters(
          actions=["1", "2"],
          rl=RLParameters(use_seq_num_diff_as_time_diff=True),
      )
      trainer = self._construct_trainer(new_params=param_copy)
      discount_tensor = trainer.compute_discount_tensor(
          batch=inp, boosted_rewards=inp.reward
      )
      self.assertEqual(discount_tensor.shape[0], self.batch_size)
      self.assertEqual(discount_tensor.shape[1], 1)
      self.assertTrue(
          torch.isclose(
              discount_tensor, torch.tensor(trainer.gamma**time_diff)
          ).all()
      )
  
      # multi_steps
      param_copy = DQNTrainerParameters(
          actions=["1", "2"], rl=RLParameters(multi_steps=steps)
      )
      trainer = self._construct_trainer(new_params=param_copy)
      discount_tensor = trainer.compute_discount_tensor(
          batch=inp, boosted_rewards=inp.reward
      )
      self.assertEqual(discount_tensor.shape[0], self.batch_size)
      self.assertEqual(discount_tensor.shape[1], 1)
      self.assertTrue(
          torch.isclose(
              discount_tensor, torch.tensor(trainer.gamma**steps)
          ).all()
      )
  
imports: import unittest
  import torch
  from reagent.core.parameters import EvaluationParameters, RLParameters
  from reagent.core.types import DiscreteDqnInput, ExtraData, FeatureData
  from reagent.evaluation.evaluation_data_page import EvaluationDataPage
  from reagent.evaluation.evaluator import get_metrics_to_score
  from reagent.models.dqn import FullyConnectedDQN
  from reagent.training.dqn_trainer import DQNTrainer
  from reagent.training.parameters import DQNTrainerParameters
  from reagent.workflow.types import RewardOptions
label_desc: Use image node_id 7 for calling the TestDQN obj's underlying member method code with example usage: obj.test_compute_discount_tensor() without return types