docstr: Test that we have the correct number of control and discount masks.
docstr_tok: ["Test","that","we","have","the","correct","number","of","control","and","discount","masks","."]
code: 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)
  
file_path: trax/trax/rl/envs/data_envs_test.py
imports: from absl.testing import absltest
  import numpy
  from trax.rl.envs import data_envs
label_desc: Use image node_id 6 for calling the SequenceDataEnvTest obj's underlying member method code with example usage: obj.test_number_of_active_masks() without return types