Skip to content

Commit f121701

Browse files
authored
[Feature] Add support of non-pickable gym env (#1615)
1 parent 5d13488 commit f121701

File tree

1 file changed

+8
-1
lines changed

1 file changed

+8
-1
lines changed

torchrl/envs/libs/gym.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1016,7 +1016,14 @@ def _build_env(
10161016
raise err
10171017
env = super()._build_env(env, pixels_only=pixels_only, from_pixels=from_pixels)
10181018
if num_envs > 0:
1019-
env = self._async_env([CloudpickleWrapper(lambda: env)] * num_envs)
1019+
try:
1020+
env = self._async_env([CloudpickleWrapper(lambda: env)] * num_envs)
1021+
except RuntimeError:
1022+
# It would fail if the environment is not pickable. In that case,
1023+
# delegating environment instantiation to each subprocess as a fallback.
1024+
env = self._async_env(
1025+
[lambda: self.lib.make(env_name, **kwargs)] * num_envs
1026+
)
10201027
self.batch_size = torch.Size([num_envs, *self.batch_size])
10211028
return env
10221029

0 commit comments

Comments
 (0)