code: def test_advantage_future_return_is_zero_at_done(self, advantage_fn):
      rewards = np.array([[1, 1, 1]], dtype=np.float32)
      returns = np.array([[3, 2, 1]], dtype=np.float32)
      values = np.array([[2, 2, 2]], dtype=np.float32)
      dones = np.array([[False, True, False]])
      discount_mask = np.ones_like(rewards)
      adv = advantage_fn(gamma=0.9, margin=1)(
          rewards, returns, values, dones, discount_mask
      )
      target_returns = values[:, :-1] + adv
      # Assert that in the "done" state the future return in the advantage is
      # zero, i.e. the return is equal to the reward.
      np.testing.assert_almost_equal(
          target_returns[0, 1], rewards[0, 1]
      )
  
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 9 for calling the AdvantagesTest obj's underlying member method code with example usage: obj.test_advantage_future_return_is_zero_at_done(advantage_fn) without return types