code: def test_gae_hardcoded(self):
      # test on small example that can be verified by hand
      rewards = np.array([[1.0, 0.0], [0.0, 0.0], [-1.0, 1.0]])
      # one of the two episodes terminated in the middle
      terminal_masks = np.array([[1.0, 1.0], [0.0, 1.0], [1.0, 1.0]])
      values = np.array(
          [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0], [1.0, 1.0]]
      )
      discount = 0.5
      gae_param = 0.25
      correct_gae = np.array(
          [[0.375, -0.5546875], [-1.0, -0.4375], [-1.5, 0.5]]
      )
      actual_gae = ppo_lib.gae_advantages(
          rewards, terminal_masks, values, discount, gae_param
      )
      np_testing.assert_allclose(actual_gae, correct_gae)
  
imports: from absl.testing import absltest
  from flax.training import train_state
  import jax
  import ml_collections
  import numpy
  import numpy.testing
  import agent
  import env_utils
  import models
  import ppo_lib
label_desc: Use image node_id 2 for calling the TestGAE obj's underlying member method code with example usage: obj.test_gae_hardcoded() without return types