Added fixes based on maintainer's feedbacks

Author: nluu175

compiler_gym/wrappers/time_limit.py

@@ -32,18 +32,8 @@ def __init__(self, env: CompilerEnv, max_episode_steps: Optional[int] = None):
         self._max_episode_steps = max_episode_steps
         self._elapsed_steps = None
 
-    # def step(self, action: ActionType, **kwargs):
-    #     assert (
-    #         self._elapsed_steps is not None
-    #     ), "Cannot call env.step() before calling reset()"
-    #     observation, reward, done, info = self.env.step(action, **kwargs)
-    #     self._elapsed_steps += 1
-    #     if self._elapsed_steps >= self._max_episode_steps:
-    #         info["TimeLimit.truncated"] = not done
-    #         done = True
-    #     return observation, reward, done, info
-
     def multistep(self, actions: Iterable[ActionType], **kwargs):
+        actions = list(actions)
         assert (
             self._elapsed_steps is not None
         ), "Cannot call env.step() before calling reset()"

tests/wrappers/time_limit_wrappers_test.py

@@ -81,20 +81,26 @@ def test_time_limit_fork(env: LlvmEnv):
         fkd.close()
 
 
-# @pytest.mark.xfail(strict=True, reason="https://github.com/facebookresearch/CompilerGym/issues/648")
 def test_time_limit(env: LlvmEnv):
     """Check CycleOverBenchmarks does not break TimeLimit"""
-    env = TimeLimit(env, max_episode_steps=1)
+    env = TimeLimit(env, max_episode_steps=3)
     env = CycleOverBenchmarks(
         env,
         benchmarks=[
             "benchmark://cbench-v1/crc32",
         ],
     )
     env.reset()
-    _, _, done, _ = env.step(0)
 
+    _, _, done, info = env.step(0)
+    assert not done, info
+
+    _, _, done, info = env.step(0)
+    assert not done, info
+
+    _, _, done, info = env.step(0)
     assert done
+    assert info["TimeLimit.truncated"], info
 
 
 if __name__ == "__main__":