def _step(self, actions):
    """Takes a step in all environments, shouldn't pre-process or record.

    Args:
      actions: (np.ndarray) with first dimension equal to the batch size.

    Returns:
      a tuple of stacked raw observations, raw rewards, dones and infos.
    """
    assert len(actions) == len(self._envs)

    observations = [None] * self.batch_size
    rewards = [None] * self.batch_size
    dones = [None] * self.batch_size
    infos = [{} for _ in range(self.batch_size)]

    def apply_step(i):
      t1 = time.time()
      observations[i], rewards[i], dones[i], infos[i] = self._envs[i].step(
          actions[i])
      t2 = time.time()
      infos[i]["__bare_env_run_time__"] = t2 - t1

    if self._parallelism > 1:
      self._pool.map(apply_step, range(self.batch_size))
    else:
      for i in range(self.batch_size):
        apply_step(i)

    # Convert each list (observations, rewards, ...) into np.array and return a
    # tuple.
    return tuple(map(_stack, [observations, rewards, dones, infos]))
Takes a step in all environments, shouldn't pre-process or record.

    Args:
      actions: (np.ndarray) with first dimension equal to the batch size.

    Returns:
      a tuple of stacked raw observations, raw rewards, dones and infos.
    
tensor2tensor/envs/gym_env_problem.py