def test_nonempty_trajectory_last_observation(self):
    """Test that last_observation is the one passed in the last extend()."""
    tr = rl_task.Trajectory(observation=123)
    for _ in range(5):
      self._extend(tr)
    self._extend(tr, new_observation=321)
    self.assertEqual(tr.last_observation, 321)
Test that last_observation is the one passed in the last extend().
trax/rl/task_test.py