From c86589b3b38563630c56cc4a041e5a99fab7662b Mon Sep 17 00:00:00 2001 From: Alexis DUBURCQ Date: Sun, 8 Oct 2023 10:57:24 +0200 Subject: [PATCH 1/3] Add support of non-pickable gym env --- torchrl/envs/libs/gym.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py index 5be494d01bb..5c37672ece1 100644 --- a/torchrl/envs/libs/gym.py +++ b/torchrl/envs/libs/gym.py @@ -1016,7 +1016,12 @@ def _build_env( raise err env = super()._build_env(env, pixels_only=pixels_only, from_pixels=from_pixels) if num_envs > 0: - env = self._async_env([CloudpickleWrapper(lambda: env)] * num_envs) + try: + env = self._async_env([CloudpickleWrapper(lambda: env)] * num_envs) + except RuntimeError: + # It would fail if the environment is not pickable. In such a case, + # delegating environment instantiation to the subprocesses as a falling. + env = self._async_env([lambda: self.lib.make(env_name, **kwargs)] * num_envs) self.batch_size = torch.Size([num_envs, *self.batch_size]) return env From 5e8c5c49e2a60138bb7a6c3a082c9c507cf40a0f Mon Sep 17 00:00:00 2001 From: Alexis DUBURCQ Date: Sun, 8 Oct 2023 19:28:21 +0200 Subject: [PATCH 2/3] Fix formatting. --- torchrl/envs/libs/gym.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py index 5c37672ece1..5edc283d9e3 100644 --- a/torchrl/envs/libs/gym.py +++ b/torchrl/envs/libs/gym.py @@ -1021,7 +1021,9 @@ def _build_env( except RuntimeError: # It would fail if the environment is not pickable. In such a case, # delegating environment instantiation to the subprocesses as a falling. - env = self._async_env([lambda: self.lib.make(env_name, **kwargs)] * num_envs) + env = self._async_env( + [lambda: self.lib.make(env_name, **kwargs)] * num_envs + ) self.batch_size = torch.Size([num_envs, *self.batch_size]) return env From 102c2a53c4f5ca05961cf48ae830dfa016416406 Mon Sep 17 00:00:00 2001 From: Alexis DUBURCQ Date: Sun, 8 Oct 2023 19:30:30 +0200 Subject: [PATCH 3/3] Fix comment. --- torchrl/envs/libs/gym.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py index 5edc283d9e3..fbec9c4f657 100644 --- a/torchrl/envs/libs/gym.py +++ b/torchrl/envs/libs/gym.py @@ -1019,8 +1019,8 @@ def _build_env( try: env = self._async_env([CloudpickleWrapper(lambda: env)] * num_envs) except RuntimeError: - # It would fail if the environment is not pickable. In such a case, - # delegating environment instantiation to the subprocesses as a falling. + # It would fail if the environment is not pickable. In that case, + # delegating environment instantiation to each subprocess as a fallback. env = self._async_env( [lambda: self.lib.make(env_name, **kwargs)] * num_envs )