尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

Gymnasium 向量环境工具函数完全指南:Space 批处理、共享内存与进程通信

Gymnasium 向量环境工具函数完全指南:Space 批处理、共享内存与进程通信 Gymnasium 向量环境工具函数完全指南Space 批处理、共享内存与进程通信【免费下载链接】GymnasiumA standard API for single-agent reinforcement learning environments, with popular reference environments and related utilities (formerly Gym)项目地址: https://gitcode.com/GitHub_Trending/gy/GymnasiumGymnasium 的gymnasium.vector子包提供了将多个环境实例并行化为向量环境的完整方案而gymnasium.vector.utils则是这套方案中承上启下的工具箱它负责把单个环境的Space批量化batching、把多个环境的样本拼接/拆分、在进程间分配共享内存以及封装环境工厂函数以支持多进程序列化。本文以 docs/api/vector/utils.md 为主线逐函数剖析其接口语义、底层实现与典型使用场景读完你既能掌握每个工具函数的签名与行为也能理解SyncVectorEnv与AsyncVectorEnv内部是如何依赖这些函数完成观测聚合与进程通信的。概述向量环境工具函数的三大板块从文档结构看gymnasium.vector.utils提供的工具被组织为三组职责Vectorizing SpacesSpace 批量化batch_space、concatenate、iterate、create_empty_array负责把单环境 Space 变为批处理 Space并在批量样本与单环境样本之间转换Shared Memory for a SpaceSpace 的共享内存create_shared_memory、read_from_shared_memory、write_to_shared_memory负责在多个进程之间共享观测数据的底层存储Miscellaneous其他杂项CloudpickleWrapper、clear_mpi_env_vars负责多进程启动时的序列化与环境变量清理。全部 10 个函数都在 gymnasium/vector/utils/init.py 中统一导出并汇总为__all__列表因此调用方式统一为from gymnasium.vector.utils import batch_space等形式。这些函数大量使用functools.singledispatch按 Space 类型Box、Discrete、MultiDiscrete、MultiBinary、Tuple、Dict、Graph、Text、Sequence、OneOf分派实现用户也可以为自定义 Space 注册新的分派实现这是理解整套 API 的关键前提。一、Vectorizing Spaces把单环境 Space 变成神经网络友好的批量结构向量环境最大的价值是为强化学习训练提供形状规整的批量数据batch从而可以直接喂给神经网络。这一组函数完成了单环境 Space ⇄ 批量 Space ⇄ 批量样本的全链路转换。文档中该板块的 4 个函数全部定义在 gymnasium/vector/utils/space_utils.py。1.1 batch_space(space, n)生成大小为 n 的批量化 Spacefrom gymnasium.vector.utils import batch_space签名batch_space(space: Space[Any], n: int 1) - Space[Any]语义把单个环境的 Space 批量化成能容纳 n 个样本的 Space。space是单个环境即向量环境中的子环境的观测/动作空间n是批量大小即向量环境中的环境数量。其返回的批量化 Space 是为神经网络优化的——也就是说它对不同 Space 类型采用了各异的批量化策略并非简单地把 n 个空间塞进一个Tuple原始 Space 类型批量化结果说明BoxBox通过np.tile在每个维度前插入长度为 n 的新轴形状从(d1, d2, ...)变为(n, d1, d2, ...)low/high 同步扩展DiscreteMultiDiscrete生成nvec全为space.n、start全为space.start的MultiDiscreteMultiDiscreteBoxlow 平铺starthigh 为start nvec - 1MultiBinaryBox形状(n,) space.shapelow0、high1TupleTuple对每个子空间递归调用batch_spaceDictDict对每个键对应的子空间递归调用batch_space保留sort_keys设置Graph/Text/Sequence/OneOf及其他Tuple回退为包含 n 个深拷贝子空间的Tuple实现见 space_utils.py。需要注意两个实现细节随机数隔离批量化后的 Space 通过deepcopy(space.np_random)复制随机数生成器而不是引用原对象。这样对原 Space 采样不会影响批量化 Space 的采样序列反之亦然源码注释明确指出如果不 deepcopy批量化后space.np_random会与batched_space.spaces[0].np_random共享同一个对象采样时就会互相干扰。对于回退到Tuple的Graph/Text/Sequence/OneOf还会额外为每个子空间重新seed一组新种子见 space_utils.py 中的new_seeds逻辑确保各子空间采样彼此独立。未注册类型的报错如果传入的不是Space实例batch_space的基函数会抛出TypeError自定义 Space 若未注册分派实现同样会触发基函数提示注册该类型后即可支持。文档给出的示例可以直观地看到批量化结果from gymnasium.spaces import Box, Dict import numpy as np space Dict({ position: Box(low0, high1, shape(3,), dtypenp.float32), velocity: Box(low0, high1, shape(2,), dtypenp.float32) }) batch_space(space, n5) # Dict(position: Box(0.0, 1.0, (5, 3), float32), velocity: Box(0.0, 1.0, (5, 2), float32))可以看到两个子空间的形状都在最前面多了一个维度 5这正是向量环境observation_space的典型形态。实际上sync_vector_env.py 和 async_vector_env.py 在初始化时正是用batch_space(self.single_action_space, self.num_envs)生成self.action_space的。batch_differing_spaces处理子空间存在差异的情况源码中还提供了文档提到的姊妹函数batch_differing_spaces(spaces)见 space_utils.py用于批量化一组类型相同但参数略有差异的 Space。例如对[Discrete(3), Discrete(5), Discrete(4), Discrete(8)]会返回MultiDiscrete([3 5 4 8])——n 个互不相同的Discrete直接摊平成一个MultiDiscrete向量。它要求所有 Space 类型一致否则抛出TypeError对Box还强制要求 dtype、low/high 的 shape 一致对Dict要求键集合一致。测试 test_space_utils.py 验证了batch_differing_spaces产出的批量样本逐个子样本仍落在原始 Space 内。在AsyncVectorEnv(..., observation_modedifferent)下正是通过该函数把各子环境不同的观测空间合并为一个批量 Space见 async_vector_env.py。1.2 concatenate(space, items, out)把多个单环境样本拼接为批量样本签名concatenate(space: Space, items: Iterable, out: tuple | Mapping | np.ndarray) - tuple | dict | np.ndarray语义space是每个样本所属的未批量化Space即向量环境中的single_action_space/single_observation_spaceitems是需要拼接的样本集合每个样本都应是space的合法元素out是输出对象通常由create_empty_array预先创建。返回拼接后的批量对象——返回值可能直接就是out本身文档在 doctest 中通过concatenate(space, items, out)原地写入预分配的out。对Box/Discrete/MultiDiscrete/MultiBinary四类基础 Space实现是np.stack(list(items), axis0, outout)见 space_utils.py即沿新轴 0 堆叠这正是给神经网络批量喂数据的标准布局。对Tuple与Dict则递归地对每个子空间/子键拼接。对Graph/Text/Sequence/OneOf等无法用定长 numpy 数组表示的 Space退化为返回tuple(items)保持原样本不做真正的拼接。文档中的示例from gymnasium.spaces import Box import numpy as np space Box(low0, high1, shape(3,), seed42, dtypenp.float32) out np.zeros((2, 3), dtypenp.float32) items [space.sample() for _ in range(2)] concatenate(space, items, out) # array([[0.77395606, 0.43887845, 0.85859793], # [0.697368 , 0.09417735, 0.97562236]], dtypefloat32)在 sync_vector_env.py 的reset方法中各子环境返回的观测列表正是通过concatenate(self.single_observation_space, self._env_obs, self._observations)被拼进预先用create_empty_array分配的缓冲区的step方法同样如此sync_vector_env.py。AsyncVectorEnv的reset_wait/step_wait也以相同模式使用concatenate汇总各 worker 进程返回的观测见 async_vector_env.py。1.3 iterate(space, items)把批量样本拆回一个个单环境样本签名iterate(space: Space[_T], items: _T) - Iterator[Any]语义space是批量化后的Spaceitems是该 Space 下的批量样本。它返回一个迭代器逐次产出批量化之前的单个样本——与concatenate构成互逆操作。对Box/MultiDiscrete/MultiBinary直接iter(items)对Dict则逐键迭代并在每次迭代时把各键的第 i 个元素重组为一个dict见 space_utils.py对Tuple若子空间全部已注册则用zip逐元素组合否则回退iter(items)并在失败时报出未注册的自定义 Space 类型。Discrete明确不支持迭代抛TypeError因为单个离散整数无法拆分为多个子样本。文档中的Dict示例可以直观看到按环境切分的效果from gymnasium.spaces import Box, Dict import numpy as np space Dict({ position: Box(low0, high1, shape(2, 3), seed42, dtypenp.float32), velocity: Box(low0, high1, shape(2, 2), seed42, dtypenp.float32) }) items space.sample() it iterate(space, items) next(it) # {position: array([0.77395606, 0.43887845, 0.85859793], dtypefloat32), # velocity: array([0.77395606, 0.43887845], dtypefloat32)} next(it) # {position: array([0.697368, 0.09417735, 0.97562236], dtypefloat32), # velocity: array([0.85859793, 0.697368], dtypefloat32)}向量环境的step就是靠它把用户传入的批量 action 拆开、逐个分发给子环境actions_iter iterate(self.action_space, actions)见 sync_vector_env.py 与 async_vector_env.py。1.4 create_empty_array(space, n, fn)预分配批量样本的容器签名create_empty_array(space: Space, n: int 1, fn: Callable np.zeros) - tuple | dict | np.ndarray | None语义创建一个可能是嵌套的、通常基于 numpy 的空数组专门与concatenate(..., outarray)配套使用——先按space的形状分配(n,) space.shape的容器再在运行时把各环境样本原地写入。space是向量环境中单个环境的观测空间n是环境数量fn是创建数组的函数常见取值包括np.zeros与np.empty。需要特别留意的是当nNone时该函数会退化为创建单个空样本文档注明如果nNone则创建space的一个空样本。分派规则为Box/Discrete/MultiDiscrete/MultiBinaryfn((n,) space.shape, dtypespace.dtype)Tuple/Dict递归调用Graph返回n个GraphInstance其中edge_links固定为形状(1, 2)的int64数组无edge_space时为NoneText返回n个由space.characters[0] * space.min_length构成的最小合法字符串SequencestackTrue时按feature_space创建否则返回 n 个空tupleOneOf返回 n 个空tuple未注册的自定义 Space返回None。文档示例n2, fnnp.zerosfrom gymnasium.spaces import Box, Dict import numpy as np space Dict({ position: Box(low0, high1, shape(3,), dtypenp.float32), velocity: Box(low0, high1, shape(2,), dtypenp.float32), }) create_empty_array(space, n2, fnnp.zeros) # {position: array([[0., 0., 0.], # [0., 0., 0.]], dtypefloat32), # velocity: array([[0., 0.], # [0., 0.]], dtypefloat32)}SyncVectorEnv在__init__中调用create_empty_array(self.single_observation_space, nself.num_envs, fnnp.zeros)初始化self._observations见 sync_vector_env.pyAsyncVectorEnv在shared_memoryFalse时同样用它分配观测缓冲区见 async_vector_env.py。1.5 组合使用batch → sample → iterate → concatenate 的闭环以上 4 个函数在设计上是可以串联使用的一整套流程。测试 test_space_utils.py 展示了完整闭环batch_space(space, n)得到批量 Space从批量 Space 采样得到批量样本iterate(batched_space, batched_sample)拆回 n 个单环境样本验证每个样本都属于原始space从原始space独立采样 n 个样本create_empty_array(space, n)预分配容器concatenate(space, samples, array)把 n 个样本拼进容器再次iterate拆回并与原始样本逐一比对使用data_equivalence断言数据等价。同文件还覆盖了确定性相同种子产生相同批量化结果且随机数互不串扰、Dict键顺序保持sort_keysFalse时batch_space不重排键见 test_space_utils.py、以及非 Space 输入抛TypeError等边界行为。二、Shared Memory for a Space多进程环境共享观测数据在AsyncVectorEnv中多个环境各自运行在独立的 worker 进程中子进程把观测写回主进程时最直接的方式是通过 Pipe 逐份传输 numpy 数组但序列化开销大。Gymnasium 的替代方案是共享内存主进程预先分配一块所有进程都能访问的内存基于multiprocessing.sharedctypes.SynchronizedArrayworker 直接写入、主进程直接读取完全绕开拷贝。这三个函数定义在 gymnasium/vector/utils/shared_memory.py。2.1 create_shared_memory(space, n, ctx)为 Space 分配跨进程共享内存签名create_shared_memory(space: Space[Any], n: int 1, ctx: ModuleType mp) - dict | tuple | SynchronizedArray语义为向量环境创建共享内存对象最终用于存放所有子环境的观测。space是单环境观测空间n是环境数量即进程数ctx是 multiprocessing 模块默认mp可传入multiprocessing.get_context(fork/spawn)得到的上下文控制进程启动方式。分配规则见 shared_memory.pyBox/Discrete/MultiDiscrete/MultiBinary计算总元素数n * prod(space.shape)booldtype 用c_bool有对应 typecode 的 dtype 直接用ctx.Array(dtype_char, size)没有 typecode 的 dtype如float16按字节数分配c_uint8数组读写时再按space.dtype重新解释原始缓冲区Tuple/Dict递归为每个子空间分配Text分配n * space.max_length个int32按字符索引存储字符串OneOf返回(int64 数组 各子空间共享内存)的元组int64数组记录每个样本选择了哪个子空间Graph与Sequence直接抛TypeError因为它们是动态形状无法预分配静态共享内存源码明确提示对于AsyncVectorEnv请禁用shared_memoryshared_memory.py。AsyncVectorEnv初始化时即调用create_shared_memory(self.single_observation_space, nself.num_envs, ctxctx)创建缓冲区见 async_vector_env.py该缓冲区随后作为_obs_buffer传入每个 worker 进程。2.2 read_from_shared_memory(space, shared_memory, n)从共享内存读出批量观测签名read_from_shared_memory(space: Space, shared_memory, n: int 1) - dict | tuple | np.ndarray语义把共享内存中的一批观测读取为可能嵌套的numpy 数组。对基础 Space通过np.frombuffer(shared_memory.get_obj(), dtypespace.dtype).reshape((n,) space.shape)实现shared_memory.py——返回的数组与共享内存共享同一块底层内存。文档中的.notes明确警告对read_from_shared_memory返回的数组做的任何修改都会同步到共享内存反之亦然如果想避免副作用必须使用np.copy。Text的读取会把int32索引序列重新映射回字符集拼成字符串元组shared_memory.pyOneOf则先读每个样本选择的子空间索引再从对应子空间的共享内存中取出该样本shared_memory.py。2.3 write_to_shared_memory(space, index, value, shared_memory)把单个环境观测写入共享内存签名write_to_shared_memory(space: Space, index: int, value: np.ndarray, shared_memory) - None语义把第index个环境index必须落在[0, num_envs)的单个观测写入共享内存。对基础 Space实现是把 value 展平后np.copyto到共享内存对应index的切片shared_memory.py。Tuple/Dict递归写入各子空间Text用gymnasium.spaces.flatten把字符串转成int32索引序列再写入OneOf则把子空间索引 值分开写入索引写进int64数组值写进对应子空间的共享内存见 shared_memory.py。在AsyncVectorEnv的 worker 循环中每个 worker 在产生新的观测后即调用write_to_shared_memory(self.single_observation_space, i, observation, shared_memory)见 async_vector_env.py 与 async_vector_env.py把结果直接写入共享内存从而省去大数组的 Pipe 传输。2.4 三函数组合与边界行为测试 test_shared_memory.py 覆盖了 create → write → read 的完整往返先batch_space(space, nnum)再create_shared_memory对每个索引write_to_shared_memory最后read_from_shared_memory读回并逐一与原始样本做data_equivalence比对。测试还对fork/spawn两种启动上下文num1与num8做了参数化验证并在当前平台不支持某启动方式时自动跳过。两个值得注意的边界无 typecode 的 dtype测试 test_shared_memory.py 专门验证了float16以及平台可用时的float128能通过按字节分配 重解释机制完整往返读取后 dtype 与 shape 均保持不变自定义 Space对未注册create_shared_memory/read_from_shared_memory/write_to_shared_memory分派实现的 Space抛CustomSpaceError对非 Space 输入抛TypeError见 test_shared_memory.py。AsyncVectorEnv会把这些错误包装成建议shared_memoryFalse的ValueErrorasync_vector_env.py。三、Miscellaneous进程启动的序列化与 MPI 环境清理这一组两个工具gymnasium/vector/utils/misc.py服务于AsyncVectorEnv多进程的启动阶段。3.1 CloudpickleWrapper用 cloudpickle 序列化环境工厂函数CloudpickleWrapper是一个可调用对象它把创建环境的工厂函数fn: Callable[[], Env]包起来并重写__getstate__/__setstate__让 Python 的 pickling 机制使用cloudpickle而非标准 pickle 来序列化该函数__getstate__调用cloudpickle.dumps(self.fn)返回字节串__setstate__用标准pickle.loads还原函数__call__无参调用self.fn()返回环境实例。为什么需要 cloudpickle因为标准pickle只能序列化模块顶层的函数而用户传给AsyncVectorEnv的env_fn往往是定义在脚本/交互环境/闭包中的函数例如lambda: gym.make(CartPole-v1)只有 cloudpickle 能保留这类函数的完整定义。在 async_vector_env.py 中每个 worker 进程都以CloudpickleWrapper(env_fn)作为ctx.Process的参数启动子进程侧还原后调用它创建自己的环境实例。注意__setstate__用的是标准库pickle——因为 cloudpickle 序列化出的字节流在还原时与标准 pickle 兼容而dumps阶段必须用 cloudpickle。3.2 clear_mpi_env_vars启动子进程前临时清理 MPI 环境变量clear_mpi_env_vars是一个上下文管理器contextlib.contextmanager用于临时删除环境变量中所有以OMPI_或PMI_前缀开头的变量退出上下文时再恢复原状misc.py。背景是from mpi4py import MPI默认会调用MPI_Init如果子进程继承了 MPI 相关的环境变量MPI 会误以为子进程与父进程一样都是 MPI 进程从而做出危险的举动例如挂起。因此AsyncVectorEnv在ctx.Process(...).start()启动所有 worker 之前会用with clear_mpi_env_vars():包裹整个进程创建循环async_vector_env.py保证子进程不会继承这些 MPI 变量无论循环中是否抛异常finally分支都会把变量恢复避免影响父进程后续行为。四、扩展能力为自定义 Space 注册工具函数整套utilsAPI 均基于functools.singledispatch构建这意味着用户可以把自己的自定义 Space 类型注册进这些函数的分派表让向量环境完整支持自定义空间。以batch_space为例注册方式为from gymnasium.vector.utils import batch_space from my_space import MySpace batch_space.register(MySpace) def _batch_my_space(space: MySpace, n: int 1): # 返回某种能容纳 n 个 MySpace 样本的空间 ...如果某个自定义 Space 没有注册各函数的基函数会给出明确错误batch_space抛TypeError而iterate、create_shared_memory、read_from_shared_memory、write_to_shared_memory抛CustomSpaceError见 space_utils.py 与 shared_memory.py。测试 test_space_utils.py 展示了未注册自定义 Space 的行为batch_space回退为Tupleconcatenate回退为tuple(items)create_empty_array返回None而iterate抛出CustomSpaceError。五、在向量环境中的真实调用链把上面的函数串起来可以完整还原向量环境的数据流以SyncVectorEnv为例见 gymnasium/vector/sync_vector_env.py初始化batch_space(single_action_space, num_envs)生成批量 action 空间batch_space或batch_differing_spaces取决于observation_mode生成批量观测空间create_empty_array预分配观测缓冲区reset各子环境返回观测后concatenate(single_observation_space, env_obs, buffer)把 n 份观测原地拼进缓冲区stepiterate(action_space, actions)把批量 action 拆成 n 份逐个传给子环境返回的观测再次concatenate聚合AsyncVectorEnv 差异额外用create_shared_memory分配共享缓冲worker 用write_to_shared_memory写入、主进程用read_from_shared_memory读取仅当shared_memoryTrue对Graph/Sequence等动态形状空间需关闭此选项进程启动依赖CloudpickleWrapper与clear_mpi_env_vars。这一调用关系同样可以在 async_vector_env.py 中逐一印证。如果想要更完整地了解VectorEnv基类与两种实现的接口可继续阅读 docs/api/vector.md、docs/api/vector/sync_vector_env.md 与 docs/api/vector/async_vector_env.md向量环境的创建入口gymnasium.make_vec与注册机制见 docs/api/registry.md。结语gymnasium.vector.utils虽然看起来只是一组小工具实则是向量环境高性能与通用性的根基batch_space/concatenate/iterate/create_empty_array定义了单环境 Space 与批量数据结构之间的标准换算create_shared_memory/read_from_shared_memory/write_to_shared_memory让多进程环境绕开昂贵的序列化传输而CloudpickleWrapper/clear_mpi_env_vars保证了 worker 进程能被正确创建。理解这 10 个函数就等于理解了SyncVectorEnv与AsyncVectorEnv内部数据流动的全部秘密当你在处理自定义 Space、混合观测空间Dict/Tuple、动态形状空间或需要跨进程共享大观测时这些工具也是你编写高性能向量化训练代码的直接抓手。【免费下载链接】GymnasiumA standard API for single-agent reinforcement learning environments, with popular reference environments and related utilities (formerly Gym)项目地址: https://gitcode.com/GitHub_Trending/gy/Gymnasium创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表