def step(self, actions, infos=None):
    """Takes a step in all environments.

    Subclasses should override _step to do the actual reset if something other
    than the default implementation is desired.

    Args:
      actions: Batch of actions.
      infos: (optional) a dictionary of keys and values, where all the values
        have the first dimension as batch_size.

    Returns:
      (preprocessed_observations, processed_rewards, dones, env_infos).
    """
    # Pre-conditions: common_preconditions, see `assert_common_preconditions`.
    #               : len(actions) == len(self._envs)
    self.assert_common_preconditions()
    assert self.batch_size == len(actions)

    observations, raw_rewards, dones, env_infos = self._step(actions)

    # Process rewards.
    raw_rewards = raw_rewards.astype(np.float32)
    processed_rewards = self.process_rewards(raw_rewards)

    # Process observations.
    processed_observations = self.process_observations(observations)

    # Record history.
    self.trajectories.step(processed_observations, raw_rewards,
                           processed_rewards, dones, actions,
                           infos=infos)

    return processed_observations, processed_rewards, dones, env_infos
Takes a step in all environments.

    Subclasses should override _step to do the actual reset if something other
    than the default implementation is desired.

    Args:
      actions: Batch of actions.
      infos: (optional) a dictionary of keys and values, where all the values
        have the first dimension as batch_size.

    Returns:
      (preprocessed_observations, processed_rewards, dones, env_infos).
    
tensor2tensor/envs/env_problem.py