def test_copy_task_short_sequence_correct_actions(self):
    """Test sequence data env on the copying task, correct replies.

    With input (x1, x2) this tests for the following sequence of
    (observations, rewards, dones, actions):
    x1                 = env.reset()
    x2,   0.0,  F, _   = env.step(ignored_action)
    eos,  0.0,  F, _   = env.step(ignored_action)
    x1,   0.0,  F, _   = env.step(x1)
    x2,   0.0,  F, _   = env.step(x2)
    eos,  1.0,  T, _   = env.step(eos)
    """
    env = data_envs.SequenceDataEnv(data_envs.copy_stream(2, n=1), 16)
    x1 = env.reset()
    x2, r0, d0, i0 = env.step(0)
    self.assertEqual(r0, 0.0)
    self.assertEqual(d0, False)
    self._assert_masks(i0, control=0, discount=0)
    eos, r1, d1, i1 = env.step(0)
    self.assertEqual(eos, 1)
    self.assertEqual(r1, 0.0)
    self.assertEqual(d1, False)
    self._assert_masks(i1, control=0, discount=0)
    y1, r2, d2, i2 = env.step(x1)
    self.assertEqual(y1, x1)
    self.assertEqual(r2, 0.0)
    self.assertEqual(d2, False)
    self._assert_masks(i2, control=1, discount=0)
    y2, r3, d3, i3 = env.step(x2)
    self.assertEqual(y2, x2)
    self.assertEqual(r3, 0.0)
    self.assertEqual(d3, False)
    self._assert_masks(i3, control=1, discount=0)
    eos2, r4, d4, i4 = env.step(1)
    self.assertEqual(eos2, 1)
    self.assertEqual(r4, 1.0)
    self.assertEqual(d4, True)
    self._assert_masks(i4, control=1, discount=1)
Test sequence data env on the copying task, correct replies.

    With input (x1, x2) this tests for the following sequence of
    (observations, rewards, dones, actions):
    x1                 = env.reset()
    x2,   0.0,  F, _   = env.step(ignored_action)
    eos,  0.0,  F, _   = env.step(ignored_action)
    x1,   0.0,  F, _   = env.step(x1)
    x2,   0.0,  F, _   = env.step(x2)
    eos,  1.0,  T, _   = env.step(eos)
    
trax/rl/envs/data_envs_test.py