def test_copy_task_max_length(self):
    """Test that sequence data env respects max_length."""
    env = data_envs.SequenceDataEnv(data_envs.copy_stream(10, n=1), 16,
                                    max_length=2)
    obs = env.reset()
    for _ in range(10):
      obs, reward, done, _ = env.step(0)
      self.assertEqual(reward, 0.0)
      self.assertEqual(done, False)
    self.assertEqual(obs, 1)  # produces EOS
    obs, reward, done, _ = env.step(7)
    self.assertEqual(obs, 7)  # repeats action
    self.assertEqual(reward, 0.0)
    self.assertEqual(done, False)
    obs, reward, done, _ = env.step(8)
    self.assertEqual(obs, 8)  # repeats action
    self.assertEqual(reward, 0.0)
    self.assertEqual(done, False)
    obs, reward, done, _ = env.step(9)
    self.assertEqual(done, True)  # exceeded max_length, stop
    self.assertEqual(obs, 1)      # produce EOS on done
    obs, reward, done, _ = env.step(10)
    self.assertEqual(done, True)  # continue producing done = True
    self.assertEqual(obs, 1)      # continue producing EOS
Test that sequence data env respects max_length.
trax/rl/envs/data_envs_test.py