def test_trajectory_suffix_observations(self):
    """Test that a trajectory suffix has the correct observations."""
    tr = rl_task.Trajectory(observation=0)
    for obs in range(1, 6):
      self._extend(tr, new_observation=obs)
    tr_suffix = tr.suffix(length=4)
    self.assertEqual([ts.observation for ts in tr_suffix.timesteps], [2, 3, 4])
    self.assertEqual(tr_suffix.last_observation, 5)
Test that a trajectory suffix has the correct observations.
trax/rl/task_test.py