fixbug(zlx): Withdraw changes in sample serial collector; In PoolEnvM…

…anager, reset returns until all envs are reset successfully

.github/workflows/envpool_test.yml

@@ -6,12 +6,11 @@ name: envpool_test
 on: [push, pull_request]
 
 jobs:
-  test_unittest:
-    runs-on: ${{ matrix.os }}
+  test_envpooltest:
+    runs-on: ubuntu-latest
     if: "!contains(github.event.head_commit.message, 'ci skip')"
     strategy:
       matrix:
-        # os: [windows-latest]
         python-version: [3.7, 3.8] # Envpool only supports python>=3.7
 
     steps:
@@ -25,6 +24,5 @@ jobs:
           python -m pip install .
           python -m pip install ".[test,k8s]"
           python -m pip install ".[envpool]"
-          # python -m pip uninstall pytest-timeouts -y
           ./ding/scripts/install-k8s-tools.sh
           make envpooltest

ding/envs/env_manager/envpool_env_manager.py

@@ -13,7 +13,7 @@
     envpool = None
 
 from ding.envs import BaseEnvTimestep
-from ding.utils import ENV_MANAGER_REGISTRY
+from ding.utils import ENV_MANAGER_REGISTRY, deep_merge_dicts
 
 
 @ENV_MANAGER_REGISTRY.register('env_pool')
@@ -47,7 +47,7 @@ def __init__(self, cfg: EasyDict) -> None:
         self._env_num = cfg.env_num
         self._batch_size = cfg.batch_size
         self._seed = cfg.seed
-        self._ready_obs = None
+        self._ready_obs = {}
         self._closed = True
 
     def launch(self) -> None:
@@ -59,11 +59,13 @@ def launch(self) -> None:
         self.reset()
 
     def reset(self) -> None:
-        self._envs.async_reset()
-        obs, _, _, _ = self._envs.recv()
-        # obs = self._envs.reset()
-        obs = obs.astype(np.float32)
-        self._ready_obs = {i: o for i, o in enumerate(obs)}
+        while True:
+            self._envs.async_reset()
+            obs, _, _, _ = self._envs.recv()
+            obs = obs.astype(np.float32)
+            self._ready_obs = deep_merge_dicts({i: o for i, o in enumerate(obs)}, self._ready_obs)
+            if len(self._ready_obs) == self._env_num:
+                break
 
     def step(self, action) -> Dict[int, namedtuple]:
         env_id = np.array(list(action.keys()))

ding/worker/collector/sample_serial_collector.py

@@ -239,10 +239,6 @@ def collect(self,
             # TODO(nyz) vectorize this for loop
             for env_id, timestep in timesteps.items():
                 with self._timer:
-                    # In async mode, this env is reset successfully just now.
-                    if self._obs_pool[env_id] is None:
-                        self._obs_pool.update({env_id: timestep.obs})
-                        continue
                     if timestep.info.get('abnormal', False):
                         # If there is an abnormal timestep, reset all the related variables(including this env).
                         # suppose there is no reset param, just reset this env