Disallow further step() calls after episode completion

scienceworld/scienceworld.py

@@ -68,6 +68,12 @@ def __init__(self, taskName: str = None, serverPath: str = None, envStepLimit: i
         # Keep track of the last step score, to calculate reward from score
         self.lastStepScore = 0
 
+        # Keep track of whether the current episode has finished (success, failure, or
+        # step limit reached), so further step() calls don't let an agent keep scoring.
+        self.isCompleted = False
+        self._lastObservation = None
+        self._lastInfos = None
+
         # Load the script
         self.taskName = taskName
         if self.taskName:
@@ -131,6 +137,9 @@ def load(self, taskName: str, variationIdx: int = 0, simplificationStr: str = ""
         # Reset last step score (used to calculate reward from current-previous score)
         self.lastStepScore = 0
 
+        # A newly loaded episode hasn't completed yet.
+        self.isCompleted = False
+
         # Keep track of whether the gold path was generated, to generate verbose error messages
         self.goldPathGenerated = generateGoldPath
 
@@ -142,6 +151,9 @@ def reset(self) -> Tuple[str, Dict[str, Any]]:
         # Reset last step score (used to calculate reward from current-previous score)
         self.lastStepScore = 0
 
+        # A freshly reset episode hasn't completed yet, so the upcoming step() isn't blocked.
+        self.isCompleted = False
+
         # Make first move
         observation, score, isCompleted, info = self.step("look around")
 
@@ -417,6 +429,17 @@ def step(self, input_str: str) -> Tuple[str, int, bool, Dict[str, Any]]:
         'moves', 'score', 'reward', 'look', 'inv', 'taskDesc', 'valid', 'variationIdx', 'taskName',
         and 'simplificationStr'.
         '''
+        # If the episode already completed (success, failure, or step limit), don't forward
+        # the action to the simulator -- this prevents an agent from continuing to act (and
+        # potentially inflating its score) after the task is over. Call reset() or load() to
+        # start a new episode.
+        if self.isCompleted:
+            logger.warning(
+                "step() was called after the episode had already completed (isCompleted=True). "
+                "Action '%s' was ignored. Call reset() or load() to start a new episode.",
+                input_str)
+            return self._lastObservation, 0, self.isCompleted, self._lastInfos
+
         observation = self.server.step(input_str)
         score = int(round(100 * self.server.getScore()))        # Convert from 0-1 to 0-100
         isCompleted = self.server.getCompleted()
@@ -449,6 +472,10 @@ def step(self, input_str: str) -> Tuple[str, int, bool, Dict[str, Any]]:
             'simplificationStr': self.simplificationStr,
         }
 
+        self.isCompleted = isCompleted
+        self._lastObservation = observation
+        self._lastInfos = infos
+
         return observation, reward, isCompleted, infos
 
     # Special actions that are "free" (consume zero time)

tests/test_scienceworld.py

@@ -136,3 +136,27 @@ def test_obj_tree():
     env.reset()
     obj_tree = env.getObjectTree()
     print(obj_tree)
+
+
+def test_step_after_completion_is_ignored():
+    env = ScienceWorldEnv("1-1", envStepLimit=1)
+    env.reset()
+
+    # Keep stepping until the step limit ends the episode.
+    done = False
+    for _ in range(10):
+        obs_done, _, done, infos_done = env.step("open door to kitchen")
+        if done:
+            break
+    assert done is True
+
+    # Further steps are ignored: the action isn't applied, and the terminal state is repeated.
+    obs_after, reward_after, done_after, infos_after = env.step("open door to kitchen")
+    assert done_after is True
+    assert reward_after == 0
+    assert obs_after == obs_done
+    assert infos_after == infos_done
+
+    # Starting a new episode clears the completed flag so step() works again.
+    env.reset()
+    assert env.isCompleted is False