Fix reset to clear previous steps

Unity-Technologies · Jul 19, 2019 · 503249b · 503249b
1 parent 4568817
commit 503249b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/ml-agents-envs/mlagents/envs/subprocess_env_manager.py b/ml-agents-envs/mlagents/envs/subprocess_env_manager.py
@@ -167,6 +167,10 @@ def step(self) -> List[StepInfo]:
     def reset(
         self, config=None, train_mode=True, custom_reset_parameters=None
     ) -> List[StepInfo]:
+        while any([ew.waiting for ew in self.env_workers]):
+            if not self.step_queue.empty():
+                step = self.step_queue.get_nowait()
+                self.env_workers[step.worker_id].waiting = False
         # First enqueue reset commands for all workers so that they reset in parallel
         for ew in self.env_workers:
             ew.send("reset", (config, train_mode, custom_reset_parameters))