[BugFix] done = done | truncated in collector (#2333)

pytorch · Jul 29, 2024 · cdcd9f3 · cdcd9f3
1 parent 0063741
commit cdcd9f3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -1121,7 +1121,10 @@ def _maybe_set_truncated(self, final_rollout):
             truncated = final_rollout["next", truncated_key]
             truncated[last_step] = True
             final_rollout["next", truncated_key] = truncated
-            final_rollout["next", _replace_last(truncated_key, "done")] = truncated
+            done = final_rollout["next", _replace_last(truncated_key, "done")]
+            final_rollout["next", _replace_last(truncated_key, "done")] = (
+                done | truncated
+            )
         return final_rollout
 
     @torch.no_grad()