Make env.close() usable after failure (finally block) (#253)

ServiceNow · Nov 13, 2024 · 3a30826 · 3a30826
1 parent 80e063b
commit 3a30826
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/browsergym/core/src/browsergym/core/env.py b/browsergym/core/src/browsergym/core/env.py
@@ -173,16 +173,22 @@ def __init__(
         self.action_space = Unicode()
 
     def close(self):
+        # stop the task
         if self.task:
-            # stop the task
             self.task.teardown()
-            # close the chat
+            self.task = None
+        # close the chat
+        if self.chat:
             self.chat.close()
-            # close the browser context
+            self.chat = None
+        # close the browser context
+        if self.context:
             self.context.close()
-            # close the browser
+            self.context = None
+        # close the browser
+        if self.browser:
             self.browser.close()
-            self.task = None
+            self.browser = None
 
     def reset(self, seed=None, *args, **kwargs):
         super().reset(seed=seed, *args, **kwargs)