text: from typing import Any, Dict, List, Optional, Tuple, Union
  
  import numpy as np
  from gym import spaces
  
  from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper
  from stable_baselines3.common.vec_env.stacked_observations import StackedDictObservations, StackedObservations
  
  
  class VecFrameStack(VecEnvWrapper):
      def __init__(
          self,
          venv: VecEnv,
          n_stack: int,
          channels_order: Optional[Union[str, Dict[str, str]]] = None,
      ):
          self.venv = venv
          self.n_stack = n_stack
          wrapped_obs_space = venv.observation_space
          if isinstance(wrapped_obs_space, spaces.Box):
              assert not isinstance(
                  channels_order, dict
              ), f"Expected None or string for channels_order but received {channels_order}"
              self.stackedobs = StackedObservations(
                  venv.num_envs, n_stack, wrapped_obs_space, channels_order
              )
          elif isinstance(wrapped_obs_space, spaces.Dict):
              self.stackedobs = StackedDictObservations(
                  venv.num_envs, n_stack, wrapped_obs_space, channels_order
              )
          else:
              raise Exception(
                  "VecFrameStack only works with gym.spaces.Box and gym.spaces.Dict observation spaces"
              )
          observation_space = self.stackedobs.stack_observation_space(wrapped_obs_space)
          VecEnvWrapper.__init__(self, venv, observation_space=observation_space)
  
      def step_wait(
          self,
      ) -> Tuple[
          Union[np.ndarray, Dict[str, np.ndarray]],
          np.ndarray,
          np.ndarray,
          List[Dict[str, Any]],
      ]:
          observations, rewards, dones, infos = self.venv.step_wait()
          observations, infos = self.stackedobs.update(observations, dones, infos)
          return observations, rewards, dones, infos
  
      def reset(self) -> Union[np.ndarray, Dict[str, np.ndarray]]:
          observation = self.venv.reset()
          observation = self.stackedobs.reset(observation)
          return observation
  
      def close(self) -> None:
          self.venv.close()
  