code: def testDiscountedReturnIsCorrectlyComputed(
      self, num_time_steps, batch_size, with_final_value
  ):
      rewards = np.random.rand(num_time_steps, batch_size).astype(
          np.float32
      )
      discounts = np.random.rand(num_time_steps, batch_size).astype(
          np.float32
      )
      final_value = (
          np.random.rand(batch_size).astype(np.float32)
          if with_final_value
          else None
      )
  
      discounted_return = value_ops.discounted_return(
          rewards=rewards, discounts=discounts, final_value=final_value
      )
  
      single_discounted_return = value_ops.discounted_return(
          rewards=rewards,
          discounts=discounts,
          final_value=final_value,
          provide_all_returns=False,
      )
  
      expected = _numpy_discounted_return(
          rewards=rewards, discounts=discounts, final_value=final_value
      )
  
      self.assertAllClose(discounted_return, expected)
      self.assertAllClose(single_discounted_return, expected[0])
  
file_path: agents/tf_agents/utils/value_ops_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.utils import value_ops
label_desc: Use image node_id 1 for calling the DiscountedReturnTest obj's underlying member method code with example usage: obj.testDiscountedReturnIsCorrectlyComputed(num_time_steps, batch_size, with_final_value) without return types