Update (base update)

[ghstack-poisoned]

Author: vmoens

test/test_env.py

@@ -13,6 +13,7 @@
 import pickle
 import random
 import re
+import time
 from collections import defaultdict
 from functools import partial
 from sys import platform
@@ -3715,26 +3716,39 @@ def test_batched_nondynamic(self, penv):
             use_buffers=True,
             mp_start_method=mp_ctx if penv is ParallelEnv else None,
         )
-        env_buffers.set_seed(0)
-        torch.manual_seed(0)
-        rollout_buffers = env_buffers.rollout(
-            20, return_contiguous=True, break_when_any_done=False
-        )
-        del env_buffers
+        try:
+            env_buffers.set_seed(0)
+            torch.manual_seed(0)
+            rollout_buffers = env_buffers.rollout(
+                20, return_contiguous=True, break_when_any_done=False
+            )
+        finally:
+            env_buffers.close(raise_if_closed=False)
+            del env_buffers
         gc.collect()
+        # Add a small delay to allow multiprocessing resource_sharer threads
+        # to fully clean up before creating the next environment. This prevents
+        # a race condition where the old resource_sharer service thread is still
+        # active when the new environment starts, causing a deadlock.
+        # See: https://bugs.python.org/issue30289
+        if penv is ParallelEnv:
+            time.sleep(0.1)
 
         env_no_buffers = penv(
             3,
             lambda: GymEnv(CARTPOLE_VERSIONED(), device=None),
             use_buffers=False,
             mp_start_method=mp_ctx if penv is ParallelEnv else None,
         )
-        env_no_buffers.set_seed(0)
-        torch.manual_seed(0)
-        rollout_no_buffers = env_no_buffers.rollout(
-            20, return_contiguous=True, break_when_any_done=False
-        )
-        del env_no_buffers
+        try:
+            env_no_buffers.set_seed(0)
+            torch.manual_seed(0)
+            rollout_no_buffers = env_no_buffers.rollout(
+                20, return_contiguous=True, break_when_any_done=False
+            )
+        finally:
+            env_no_buffers.close(raise_if_closed=False)
+            del env_no_buffers
         gc.collect()
         assert_allclose_td(rollout_buffers, rollout_no_buffers)
 

test/test_libs.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import collections
+import copy
 import functools
 import gc
 import importlib.util
@@ -2811,14 +2812,27 @@ def test_vmas_seeding(self, scenario_name):
         final_seed = []
         tdreset = []
         tdrollout = []
-        for _ in range(2):
-            env = VmasEnv(
+        rollout_length = 10
+
+        def create_env():
+            return VmasEnv(
                 scenario=scenario_name,
                 num_envs=4,
             )
+
+        env = create_env()
+        td_actions = [env.action_spec.rand() for _ in range(rollout_length)]
+
+        for _ in range(2):
+            env = create_env()
+            td_actions_buffer = copy.deepcopy(td_actions)
+
+            def policy(td, actions=td_actions_buffer):
+                return actions.pop(0)
+
             final_seed.append(env.set_seed(0))
             tdreset.append(env.reset())
-            tdrollout.append(env.rollout(max_steps=10))
+            tdrollout.append(env.rollout(max_steps=rollout_length, policy=policy))
             env.close()
             del env
         assert final_seed[0] == final_seed[1]

torchrl/data/datasets/d4rl.py

@@ -279,6 +279,7 @@ def _get_dataset_direct(self, name, env_kwargs):
         # so we need to ensure we're using the gym backend
         with set_gym_backend("gym"):
             import gym
+
             env = GymWrapper(gym.make(name))
         with tempfile.TemporaryDirectory() as tmpdir:
             os.environ["D4RL_DATASET_DIR"] = tmpdir
@@ -358,6 +359,7 @@ def _get_dataset_from_env(self, name, env_kwargs):
         # so we need to ensure we're using the gym backend
         with set_gym_backend("gym"), tempfile.TemporaryDirectory() as tmpdir:
             import gym
+
             os.environ["D4RL_DATASET_DIR"] = tmpdir
             env = GymWrapper(gym.make(name))
             dataset = make_tensordict(