env_wrappers.py

向量环境与自动 reset 接口

原文件 · 86 行 · 3103 字节

envs/delivcoal/env_wrappers.py
SHA-256: 863e088540942fb09b21a0dab9f05669920667344c82d2a1927ee598e004843d

1"""2# @Time    : 2021/7/1 8:44 上午3# @Author  : hezhiqiang014# @Email   : hezhiqiang01@baidu.com5# @File    : env_wrappers.py6Modified from OpenAI Baselines code to work with multi-agent envs7"""89import numpy as np101112def _stack_dict_obs(obs_list):13    """将 list[dict[str, ndarray]] 堆叠为 dict[str, ndarray]。"""14    keys = obs_list[0].keys()15    return {k: np.array([o[k] for o in obs_list]) for k in keys}161718# single env19class DummyVecEnv():20    def __init__(self, env_fns):21        self.envs = [fn() for fn in env_fns]22        env = self.envs[0]23        self.num_envs = len(env_fns)24        self.observation_space = env.observation_space25        self.share_observation_space = env.share_observation_space26        self.action_space = env.action_space27        self.actions = []2829    def step(self, actions, episode=-1):30        """31        Step the environments synchronously.32        This is available for backwards compatibility.33        """34        self.step_async(actions)35        return self.step_wait(episode)3637    def step_async(self, actions):38        self.actions = actions3940    def step_wait(self, episode):41        # ① 与所有子环境交互42        results = [env.step(a) for (a, env) in zip(self.actions, self.envs)]43        # results = [(obs, rew, done, info), ...]44        obs, rews, dones, infos, active_mask, bad_mask, available_actions = map(list, zip(*results))   # 用 list 保持元素可变4546        # ② 逐线程检查 done→reset47        for i, done in enumerate(dones):48            done_flag = done if isinstance(done, (bool, np.bool_)) else np.all(done)49            if done_flag:50                obs_i, active_i, available_i = self.envs[i].reset(episode)51                obs[i]   = obs_i52                active_mask[i] = active_i53                available_actions[i] = available_i5455        # ③ 转回便于上层使用的格式56        obs  = _stack_dict_obs(obs)           # dict[str, ndarray]57        rews = np.array(rews)58        dones = np.array(dones)59        active_mask = np.array(active_mask)60        bad_mask = np.array(bad_mask)61        available_actions = np.array(available_actions)62        # infos 保持 list[dict] —— Runner 需要这样63        return obs, rews, dones, infos, active_mask, bad_mask, available_actions646566    def reset(self, episode):67        result = [env.reset(episode) for env in self.envs]68        obs_list, active_masks, available_actions = zip(*result)        # unzip69        obs = _stack_dict_obs(obs_list)                                 # dict[str, ndarray]70        return obs, np.stack(active_masks), np.stack(available_actions)7172    def get_global_state(self):73        return _stack_dict_obs([env.get_global_state() for env in self.envs])7475    def close(self):76        for env in self.envs:77            env.close()7879    def render(self, mode="human"):80        if mode == "rgb_array":81            return np.array([env.render(mode=mode) for env in self.envs])82        elif mode == "human":83            for env in self.envs:84                env.render(mode=mode)85        else:86            raise NotImplementedError