Fix implementation of inheritace from ObservationWrapper and RewardWr…

…apper

examples/llvm_rl/wrappers.py

@@ -34,7 +34,7 @@ def __init__(
         self.max = max
         self.leakiness_factor = leakiness_factor
 
-    def reward(self, reward: float) -> float:
+    def convert_reward(self, reward: float) -> float:
         if reward > self.max:
             return self.max + (reward - self.max) * self.leakiness_factor
         elif reward < self.min:
@@ -65,7 +65,7 @@ def __init__(self, env: CompilerEnv):
             dtype=np.float32,
         )
 
-    def observation(self, observation):
+    def convert_observation(self, observation):
         if observation[self.TotalInsts_index] <= 0:
             return np.zeros(observation.shape, dtype=np.float32)
         return np.clip(

tests/wrappers/core_wrappers_test.py

@@ -8,26 +8,23 @@
 
 from compiler_gym.datasets import Datasets
 from compiler_gym.envs.llvm import LlvmEnv
-from compiler_gym.wrappers import (
-    ActionWrapper,
-    CompilerEnvWrapper,
-    ObservationWrapper,
-    RewardWrapper,
-)
+from compiler_gym.wrappers import ActionWrapper, CompilerEnvWrapper
+from compiler_gym.wrappers import ObservationWrapper as CoreObservationWrapper
+from compiler_gym.wrappers import RewardWrapper as CoreRewardWrapper
 from tests.test_main import main
 
 pytest_plugins = ["tests.pytest_plugins.llvm"]
 
 
-class ObservationDummyWrapper(ObservationWrapper):
+class ObservationWrapper(CoreObservationWrapper):
     def __init__(self, env):
         super().__init__(env)
 
     def convert_observation(self, observation):
         return observation
 
 
-class RewardDummyWrapper(RewardWrapper):
+class RewardWrapper(CoreRewardWrapper):
     def __init__(self, env):
         super().__init__(env)
 
@@ -40,8 +37,8 @@ def convert_reward(self, reward):
     params=[
         ActionWrapper,
         CompilerEnvWrapper,
-        ObservationDummyWrapper,
-        RewardDummyWrapper,
+        ObservationWrapper,
+        RewardWrapper,
     ],
 )
 def wrapper_type(request):