Fix failing tests

compiler_gym/bin/service.py

@@ -256,6 +256,10 @@ def print_service_capabilities(env: CompilerEnv):
                 headers=("Action",),
             )
             print(table)
+        else:
+            raise NotImplementedError(
+                "Only Commandline and NamedDiscrete are supported."
+            )
 
 
 def main(argv):

compiler_gym/envs/compiler_env.py

@@ -11,7 +11,7 @@
 from math import isclose
 from pathlib import Path
 from time import time
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import gym
 import numpy as np
@@ -1082,9 +1082,7 @@ def step(  # pylint: disable=arguments-differ
                 category=DeprecationWarning,
             )
             reward_spaces = rewards
-        return self._multistep(
-            self.raw_step, [action], observation_spaces, reward_spaces
-        )
+        return self.multistep([action], observation_spaces, reward_spaces)
 
     def multistep(
         self,
@@ -1129,20 +1127,7 @@ def multistep(
                 category=DeprecationWarning,
             )
             reward_spaces = rewards
-        return self._multistep(
-            self.raw_step, list(actions), observation_spaces, reward_spaces
-        )
 
-    def _multistep(
-        self,
-        raw_step: Callable[
-            [Iterable[ActionType], Iterable[ObservationSpaceSpec], Iterable[Reward]],
-            StepType,
-        ],
-        actions: Iterable[ActionType],
-        observation_spaces: Optional[Iterable[Union[str, ObservationSpaceSpec]]],
-        reward_spaces: Optional[Iterable[Union[str, Reward]]],
-    ) -> StepType:
         # Coerce observation spaces into a list of ObservationSpaceSpec instances.
         if observation_spaces:
             observation_spaces_to_compute: List[ObservationSpaceSpec] = [
@@ -1170,7 +1155,7 @@ def _multistep(
             reward_spaces_to_compute: List[Reward] = []
 
         # Perform the underlying environment step.
-        observation_values, reward_values, done, info = raw_step(
+        observation_values, reward_values, done, info = self.raw_step(
             actions, observation_spaces_to_compute, reward_spaces_to_compute
         )
 

examples/llvm_rl/wrappers.py

@@ -135,7 +135,7 @@ def multistep(
     ):
         for a in actions:
             self.histogram[a] += self.increment
-        return self.env.multistep(actions, **kwargs)
+        return super().multistep(actions, **kwargs)
 
     def observation(self, observation):
         return np.concatenate((observation, self.histogram)).astype(

examples/sensitivity_analysis/action_sensitivity_analysis.py

@@ -116,7 +116,7 @@ def run_one_trial(
     num_warmup_steps = random.randint(0, max_warmup_steps)
     warmup_actions = [env.action_space.sample() for _ in range(num_warmup_steps)]
     env.reward_space = reward_space
-    _, _, done, _ = env.step(warmup_actions)
+    _, _, done, _ = env.multistep(warmup_actions)
     if done:
         return None
     _, (reward,), done, _ = env.step(action, reward_spaces=[reward_space])