[BugFix] Fix async gym when all reset (pytorch#2144)

ai4co · Apr 30, 2024 · 711a4ee · 711a4ee
1 parent 7b9305d
commit 711a4ee
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/examples/envs/gym-async-info-reader.py b/examples/envs/gym-async-info-reader.py
@@ -15,6 +15,7 @@
 parser = argparse.ArgumentParser()
 parser.add_argument("--use_wrapper", action="store_true")
 
+
 # Create the dummy environment
 class CustomEnv(gym.Env):
     def __init__(self, render_mode=None):
@@ -71,6 +72,7 @@ def step(self, action):
     # Create an info reader: this object will read the info and write its content to the tensordict
     def reader(info, tensordict):
         return tensordict.set("field1", np.stack(info["field1"]))
+
     env.set_info_dict_reader(info_dict_reader=reader)
 
     # Print the info readers (there should be 2: one to read the terminal states and another to read the 'field1')

diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -2249,7 +2249,7 @@ class MultiaSyncDataCollector(_MultiDataCollector):
 
         See https://docs.python.org/3/library/multiprocessing.html for more info.
 
-        Examples:
+    Examples:
         >>> from torchrl.envs.libs.gym import GymEnv
         >>> from tensordict.nn import TensorDictModule
         >>> from torch import nn

diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -6996,7 +6996,7 @@ def _reset(
             )
         # if not reset.any(), we don't need to do anything.
         # if reset.all(), we don't either (bc GymWrapper will call a plain reset).
-        if reset is not None and reset.any() and not reset.all():
+        if reset is not None and reset.any():
             saved_next = self._memo["saved_next"]
             # reset = reset.view(tensordict.shape)
             # we have a data container from the previous call to step