need to break the loop later when action is None (#155)

Author: recursix

browsergym/experiments/src/browsergym/experiments/loop.py

@@ -218,9 +218,7 @@ def run(self):
                 logger.debug(f"Agent chose action:\n {action}")
 
                 if action is None:
-                    logger.debug(f"Agent returned None action. Ending episode.")
                     step_info.truncated = True
-                    break
 
                 step_info.save_step_info(self.exp_dir)
                 logger.debug(f"Step info saved.")
@@ -230,6 +228,11 @@ def run(self):
 
                 step_info = StepInfo(step=step_info.step + 1)
                 episode_info.append(step_info)
+
+                if action is None:
+                    logger.debug(f"Agent returned None action. Ending episode.")
+                    break
+
                 logger.debug(f"Sending action to environment.")
                 step_info.from_step(env, action, obs_preprocessor=agent.obs_preprocessor)
                 logger.debug(f"Environment stepped.")