Repository: facebookresearch/CompilerGym

Issue:
Logical Failure when combing TimeLimit Wrapper with IterateOverBenchmarks
## 🐛 Bug

If an environment is first wrapper with TimeLimit before IterateOverBenchmarks, it will not return "done" as True. 

## To Reproduce

Steps to reproduce the behavior:

```
env = TimeLimit(env, step_limit) 
env = CycleOverBenchmarks(env, benchmarks) 
_, done, _, _ = env.reset()
while not done:
    _, done, _, _ = env.step(0) 
```
This will not finish. However, if the TimeLimit happens after the Cycle, it has normal behavior. 


## Additional context

Assign it to me, I will fix it when I got time. 


Discussion:
Good catch @JD-ETH and thanks for volunteering to patch this! 🙂 

Cheers,
Chris
Hi. I tried to reproduce the behavior but it did not work. I found that there are something wrong in the code snippets that @JD-ETH provided.

For `env.step(0)`, `done` should be the third variable (currently second one in the provided code). Also, `done` is automatically assigned `True` as the beginning so the loop will not execute.

The way I reproduced the code is as follows:
```
env.reset()
i = 0
done = False
while not done:
    _, _, done, _ = env.step(0)
```
This code snippet should reproduce the exact behavior for this issue.

Resolving patch:
diff --git a/compiler_gym/wrappers/time_limit.py b/compiler_gym/wrappers/time_limit.py
--- a/compiler_gym/wrappers/time_limit.py
+++ b/compiler_gym/wrappers/time_limit.py
@@ -2,7 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-from typing import Optional
+from typing import Iterable, Optional
 
 from compiler_gym.envs import CompilerEnv
 from compiler_gym.util.gym_type_hints import ActionType
@@ -32,12 +32,13 @@ 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):
+    def multistep(self, actions: Iterable[ActionType], **kwargs):
+        actions = list(actions)
         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
+        observation, reward, done, info = self.env.multistep(actions, **kwargs)
+        self._elapsed_steps += len(actions)
         if self._elapsed_steps >= self._max_episode_steps:
             info["TimeLimit.truncated"] = not done
             done = True


Test patch:
diff --git a/tests/wrappers/time_limit_wrappers_test.py b/tests/wrappers/time_limit_wrappers_test.py
--- a/tests/wrappers/time_limit_wrappers_test.py
+++ b/tests/wrappers/time_limit_wrappers_test.py
@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 """Unit tests for //compiler_gym/wrappers."""
 from compiler_gym.envs.llvm import LlvmEnv
-from compiler_gym.wrappers import TimeLimit
+from compiler_gym.wrappers import CycleOverBenchmarks, TimeLimit
 from tests.test_main import main
 
 pytest_plugins = ["tests.pytest_plugins.llvm"]
@@ -81,5 +81,27 @@ def test_time_limit_fork(env: LlvmEnv):
         fkd.close()
 
 
+def test_time_limit(env: LlvmEnv):
+    """Check CycleOverBenchmarks does not break TimeLimit"""
+    env = TimeLimit(env, max_episode_steps=3)
+    env = CycleOverBenchmarks(
+        env,
+        benchmarks=[
+            "benchmark://cbench-v1/crc32",
+        ],
+    )
+    env.reset()
+
+    _, _, 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__":
     main()