def test_number_of_active_masks(self):
    """Test that we have the correct number of control and discount masks."""
    n_input_seqs = 3
    n_output_seqs = 2
    input_len = 4
    output_len = 5

    def data_stream():
      i = 2 * np.ones(input_len)
      o = np.zeros(output_len)
      while True:
        yield (i, o, i, o, i)  # 3 input, 2 output sequences.

    env = data_envs.SequenceDataEnv(data_stream(), 16, max_length=output_len)
    env.reset()

    n_discount = 0
    n_control = 0
    n_steps = 0
    done = False
    while not done:
      (_, _, done, info) = env.step(action=0)
      n_discount += info['discount_mask']
      n_control += info['control_mask']
      n_steps += 1

    # One discount_mask=1 per output sequence.
    self.assertEqual(n_discount, n_output_seqs)
    # One control_mask=1 per output token, including EOS, because it's also
    # controlled by the agent.
    self.assertEqual(n_control, (output_len + 1) * n_output_seqs)
    # One control_mask=0 per input token, excluding EOS, because when the env
    # emits it, control transfers to the agent immediately.
    self.assertEqual(n_steps - n_control, input_len * n_input_seqs)
Test that we have the correct number of control and discount masks.
trax/rl/envs/data_envs_test.py