code: def test_future_return_is_zero_iff_discount_mask_is_on(
      self, advantage_fn
  ):
      # (... when gamma=0)
      rewards = np.array([[1, 2, 3, 4]], dtype=np.float32)
      values = np.array([[5, 6, 7, 8]], dtype=np.float32)
      dones = np.zeros_like(rewards, dtype=bool)
      discount_mask = np.array([[1, 0, 1, 0]], dtype=bool)
      gammas = advantages.mask_discount(0.0, discount_mask)
      returns = advantages.discounted_returns(rewards, gammas)
      adv = advantage_fn(gamma=0.0, margin=1)(
          rewards, returns, values, dones, discount_mask
      )
      target_returns = values[:, :-1] + adv
      # Assert that in the states with discount_mask on the future return in the
      # advantage is zero, i.e. the return is equal to the reward.
      rewards = rewards[:, :-1]
      discount_mask = discount_mask[:, :-1]
      np.testing.assert_almost_equal(
          target_returns[discount_mask], rewards[discount_mask]
      )
      # Assert the converse.
      with np.testing.assert_raises(AssertionError):
          np.testing.assert_almost_equal(
              target_returns[~discount_mask], rewards[~discount_mask]
          )
  
file_path: trax/trax/rl/advantages_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  from trax.rl import advantages
label_desc: Use image node_id 11 for calling the AdvantagesTest obj's underlying member method code with example usage: obj.test_future_return_is_zero_iff_discount_mask_is_on(advantage_fn) without return types