RL gym 环境（3）—— 环境向量化（批量训练）

• 本文介绍如何在 gym 套件中训练向量化的环境，所谓 “向量化”，可以理解为把所有环境的 observation、action、reward、info 等所有信息都合在一起，拼成一个 “环境向量”，从而把多个独立环境的训练过程组织在一起。agent 和这些环境的交互可以是串行的，也可以是并行的。本文参考自官方文档 Vectorising your environments
• 请先看之前关于 gym 套件的基础介绍
  1. RL gym 环境（1）—— 安装和基础使用
  2. RL gym 环境（2）—— 自定义环境

---

1. 两种向量化环境

• 向量化环境Vectorized environments 将同一环境的多个独立副本组织在一起运行的环境，它输入一批动作，同时返回一批观察结果。环境向量化技术在训练时非常有用
• Gym 提供了两种类型的向量化环境
  1. 同步向量化（顺序向量化）环境 gym.vector.SyncVectorEnv：这里环境不同副本是顺序执行的
  2. 异步向量化（并行向量化）环境 gym.vector.AsyncVectorEnv：这里使用 python 的多进程机制 multiprocessing 并行地执行环境的不同副本，每个环境副本都运行在一个独立的进程中

2. 创建向量化环境

2.1 环境副本完全相同

• 如果只想在某个特定环境训练，那么可以使用该环境的相同副本组成向量化环境，这常用于加速训练
• 可以用 gym.vector.make 方法创建这种向量化环境，原型如下

  gym.vector.make(
      id: str,
      num_envs: int = 1,
      asynchronous: bool = True,
      wrappers: Union[<built-in function callable>, List[<built-in function callable>], NoneType] = None,
      disable_env_checker: Union[bool, NoneType] = None,
      **kwargs,
  ) -> gym.vector.vector_env.VectorEnv
  

  1. num_envs 指定被组织的环境副本数量
  2. asynchronous 指定是否异步交互（设为 True 异步并行）
  3. wrappers 指定各个环境副本使用的包装，这里要么都加要么都不加
  4. disable_env_checker 指定是否对第一个环境副本进行 gym 规范性检查（设为 None 或 True 不检查）
  5. kwargs 代表环境副本的自身参数
• 以前文自定义的 Cliff Walking 环境为例，可以如下将其向量化

  import gym
  #envs = gym.vector.make('CartPole-v1', num_envs=3, disable_env_checker=None)
  envs = gym.vector.make('MyGymExamples:MyGymExamples/CliffWalkingEnv-v0',
                          num_envs=3, 
                          disable_env_checker=False,
                          render_mode='rgb_array', 	# 从这开始为环境的自身参数
                          map_size=(4,12), 
                          pix_square_size=30)
  
  observations, infos = envs.reset()
  
  print('observations: ', observations)
  print('infos: ', infos)
  

  observations:  OrderedDict([('agent', array([[0, 3],
         [0, 3],
         [0, 3]], dtype=int64)), ('target', array([[11,  3],
         [11,  3],
         [11,  3]], dtype=int64))])
  infos:  {'distance': array([11., 11., 11.]), '_distance': array([ True,  True,  True])}
  

2.2 指定环境副本参数

• 2.1 节的方法在以下三种情况不适用
  1. 被组织的环境副本具有不同的参数 （比如具有不同重力 g 值的 “Pendulum-v0” 环境）
  2. 被组织的环境没有注册到 gym
  3. 在 一些（但不是全部） 环境副本上使用包装
• 这时可以用以下方式

  # 异步（并行）
  asyn_env = gym.vector.AsyncVectorEnv([
      lambda: gym.make("Pendulum-v0", g=9.81),
      lambda: gym.make("Pendulum-v0", g=1.62)
  ])
  
  # 同步（顺序）
  sync_env = gym.vector.SyncVectorEnv([
      lambda: gym.make("Pendulum-v0", g=9.81),
      lambda: gym.make("Pendulum-v0", g=1.62)
  ])
  

  需要注意的是，如果用 gym.vector.AsyncVectorEnv 创建并行训练的一组环境，由于 python 的多进程机制性质，应将其放在 if __name__ == "__main__": 之后

3. 使用向量化环境

3.1 基础用法

• 向量化环境的使用和普通环境几乎完全相同，只是原先的所有变量维度都进行了扩展，见下例

  env = gym.vector.make('MyGymExamples:MyGymExamples/CliffWalkingEnv-v0', 
                          num_envs=3, 
                          disable_env_checker=False,
                          render_mode='rgb_array', 
                          map_size=(4,12), 
                          pix_square_size=30)
  observation, info = env.reset()
  

  >>> observation
  OrderedDict([('agent',
                array([[0, 3],
                       [0, 3],
                       [0, 3]], dtype=int64)),
               ('target',
                array([[11,  3],
                       [11,  3],
                       [11,  3]], dtype=int64))])
                       
  >>> info
  {'distance': array([11., 11., 11.]), '_distance': array([ True,  True,  True])}
  

  从此可见

  1. observation 和 info 内容的维度都扩展了，
  2. info 字典中每个 key 都会多一个 _key 字段指示各个环境副本返回的该 key 字段是否有数据，其值由各个环境副本是否终止决定

• 接下来运行一步

  action_direction = {'noop': 0, 'right': 1, 'down': 2, 'left': 3, 'up': 4}
  observation, reward, terminated, truncated, info = env.step(np.array([action_direction['up'], 			# 这会向上移动一格
  																		action_direction['down'],		# 这会被地图下界挡住停在起点 
  																		action_direction['right']]))	# 这会落入悬崖   
  

  >>> observation
  OrderedDict([('agent',
                array([[0, 2],
                       [0, 3],
                       [0, 3]], dtype=int64)),
               ('target',
                array([[11,  3],
                       [11,  3],
                       [11,  3]], dtype=int64))])
                       
  >>> reward
  array([  -1,   -1, -100])
  
  >>> terminated
  array([False, False, False])
  
  >>> truncated
  array([False, False,  True])
  
  >>> info
  {'distance': array([12., 11., 11.]),
   '_distance': array([ True,  True,  True]),
   'final_observation': array([None, None, {'agent': array([1, 3]), 'target': array([11,  3])}],
         dtype=object),
   '_final_observation': array([False, False,  True]),
   'final_info': array([None, None, {'distance': 10.0}], dtype=object),
   '_final_info': array([False, False,  True])}
  

  从此可见

  1. 输入的 action、返回的 reward, terminated, truncated 等变量都进行了维度扩展
  2. 到达终止状态的子环境副本（由 terminated 或 truncated 指示）会自动 reset
  3. env.step 返回的 info 还会多几个 final_observation、final_info 相关的字段，指出了被 reset 的环境副本 reset 前对终止状态的观测

3.2 观测空间 & 动作空间

• 向量化环境和普通环境完全类似，它也有属于 gym.space 子类的 observation_space 和 action_space，这些空间是自动从被组织的环境副本推断出来的

  >>> envs = gym.vector.make("CartPole-v1", num_envs=3)
  >>> envs.observation_space
  Box([[-4.8 ...]], [[4.8 ...]], (3, 4), float32)
  
  >>> envs.action_space
  MultiDiscrete([2 2 2])
  

• 必须保证所有被组织子环境的观测和行动空间相同，否则无法正确推断

  >>> envs = gym.vector.AsyncVectorEnv([
  ...    lambda: gym.make("CartPole-v1"),
  ...    lambda: gym.make("MountainCar-v0")
  ...])
  RuntimeError: Some environments have an observation space different from `Box([-4.8 ...], [4.8 ...], (4,), float32)`. 
  In order to batch observations, the observation spaces from all environments must be equal.
  

• 对于向量化环境，可以通过 VectorEnv.single_observation_space 和 VectorEnv.single_action_space 得到其子环境副本的观测和动作空间，常用这个来指定策略模型的一些参数尺寸

  >>> envs = gym.vector.make("CartPole-v1", num_envs=3)
  >>> envs.single_observation_space
  Box([-4.8 ...], [4.8 ...], (4,), float32)
  
  >>> envs.single_action_space
  Discrete(2)
  

3.3 共享内存

• gym.vector.AsyncVectorEnv 类型的并行向量化环境会在独立进程中运行每个环境副本，每次调用 AsyncVectorEnv.reset() 或AsyncVectorEnv.step() 时，所有并行环境的 observation 结果都会发送回主进程，这种进程间数据传输成本很高，对于高维 observation 此问题尤其明显
• gym.vector.AsyncVectorEnv 默认使用进程共享内存方法 (shared_memory=True) 尽量减少进程间数据传输成本，这可以增加向量化环境的吞吐量（throughout）

3.4 异常处理

• 任何子环境副本中引发的异常都会在向量化环境中重新引发，这样我们可以自己选择如何处理这些异常，如下

  class ErrorEnv(gym.Env):
      observation_space = gym.spaces.Box(-1., 1., (2,), np.float32)
      action_space = gym.spaces.Discrete(2)
  
      def reset(self):
          return np.zeros((2,), dtype=np.float32), {}
  
      def step(self, action):
          if action == 1:
              raise ValueError("An error occurred.")
          observation = self.observation_space.sample()
          return (observation, 0., False, {})
  

  >>> envs = gym.vector.AsyncVectorEnv([lambda: ErrorEnv()] * 3)
  >>> observations, infos = envs.reset()
  >>> observations, rewards, dones, infos = envs.step(np.array([0, 0, 1]))
  ERROR: Received the following error from Worker-2: ValueError: An error occurred.
  ERROR: Shutting down Worker-2.
  ERROR: Raising the last exception back to the main process.
  ValueError: An error occurred.
  

4. 实验：通过异步向量化提高训练效率

• 本节对比普通环境 gym.Env、同步向量化环境 gym.vector.SyncVectorEnv 和异步向量化环境 gym.vector.AsyncVectorEnv 的运行速度，注意基础环境每一步交互越慢（即 .step() 用时越久），将其组织起来并行计算的效率提升将会越高，因此我们首先定义一个单步交互很慢的环境

  class SlowEnv(gym.Env):
  	# 随便定义观测和动作空间
      observation_space = gym.spaces.Dict({
          "position": gym.spaces.Box(-1., 1., (3,), np.float32),
          "velocity": gym.spaces.Box(-1., 1., (2,), np.float32)
      })
      action_space = gym.spaces.Dict({
          "fire": gym.spaces.Discrete(2),
          "jump": gym.spaces.Discrete(2),
          "acceleration": gym.spaces.Box(-1., 1., (2,), np.float32)
      })
  	
      def reset(self):
          return self.observation_space.sample(), {}
  
      def step(self, action):
          i = 0
          for _ in range(500000): i+= 1 # make it slow
          observation = self.observation_space.sample()
          return (observation, 0., False, False, {})
  

• 创建普通环境和两种向量化环境

  env = SlowEnv()
  asyn_envs = gym.vector.AsyncVectorEnv([
                  lambda: SlowEnv(),
                  lambda: SlowEnv(),
                  lambda: SlowEnv(),
              ])
  sync_envs = gym.vector.SyncVectorEnv([
                  lambda: SlowEnv(),
                  lambda: SlowEnv(),
                  lambda: SlowEnv(),
              ])
  

• 强化学习训练主要的耗时都在环境交互上，我们现在对比三种环境的交互速度，即对比 .step() 速度。可以很方便地使用 jupyter notebook 的 %timeit 魔法方法进行这种计时比较，结果如下

  >>> %timeit -n 100 -r 2 env.step(env.action_space.sample())
  26.5 ms ± 831 µs per loop (mean ± std. dev. of 2 runs, 100 loops each)
  
  >>> %timeit -n 100 -r 2 asyn_envs.step(asyn_envs.action_space.sample())
  30.4 ms ± 177 µs per loop (mean ± std. dev. of 2 runs, 100 loops each)
  
  >>> %timeit -n 100 -r 2 sync_envs.step(sync_envs.action_space.sample())
  78.9 ms ± 136 µs per loop (mean ± std. dev. of 2 runs, 100 loops each)
  

  可见异步向量化的三个环境交互耗时只比单一环境多一点，而同步量化的三个环境交互耗时约为单一环境的三倍，使用并行计算可以大幅提高训练效率。关于 %timeit 魔法方法可以参考 Jupyter Notebook %timeit 功能详解 Python 代码执行时间