From aa7b91333e1970fa84346fd1c690427a30eab03b Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Mon, 30 Dec 2019 12:01:37 +0100 Subject: [PATCH] Add seeding for subproc vecenv --- torchy_baselines/common/vec_env/subproc_vec_env.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/torchy_baselines/common/vec_env/subproc_vec_env.py b/torchy_baselines/common/vec_env/subproc_vec_env.py index b00e465..f89cbc6 100644 --- a/torchy_baselines/common/vec_env/subproc_vec_env.py +++ b/torchy_baselines/common/vec_env/subproc_vec_env.py @@ -188,6 +188,17 @@ class SubprocVecEnv(VecEnv): for remote in target_remotes: remote.recv() + def seed(self, seed, indices=None): + """ + :param seed: (int or [int]) + :param indices: ([int]) + """ + indices = self._get_indices(indices) + if not hasattr(seed, 'len'): + seed = [seed] * len(indices) + assert len(seed) == len(indices) + return [self.env_method('seed', seed[i], indices=i) for i in indices] + def env_method(self, method_name, *method_args, **method_kwargs): """Call instance methods of vectorized environments.""" indices = method_kwargs.get('indices')