def step(self, action):
    """Single step of the environment when policy took `action`."""
    cur_sequence_id, cur_token_id = self._cur_position
    if cur_sequence_id >= len(self._cur_sequence):
      return np.array(self._eos, dtype=np.int32), 0.0, True, None

    # Emit the control mask on the output.
    control_mask = int(not self._on_input)

    if self._on_input:
      self._response = []
      if cur_token_id + 1 < len(self._cur_sequence[cur_sequence_id]):
        self._cur_position = (cur_sequence_id, cur_token_id + 1)
        done = False
      else:
        self._cur_position = (cur_sequence_id + 1, 0)
        done = cur_sequence_id + 1 >= len(self._cur_sequence)
      reward = 0.0
      discount_mask = 0

    else:
      self._response.append(action)
      if action == self._eos or len(self._response) > self._max_length:
        self._cur_position = (cur_sequence_id + 1, 0)
        reward = self._metric(
            self._response[:-1], self._cur_sequence[cur_sequence_id])
        done = cur_sequence_id + 1 >= len(self._cur_sequence)
        # Emit the discount mask on the last token of each action.
        discount_mask = 1
      else:
        reward = 0.0
        done = False
        discount_mask = 0

    info = {'control_mask': control_mask, 'discount_mask': discount_mask}
    return self.observation, reward, done, info
Single step of the environment when policy took `action`.
trax/rl/envs/data_envs.py