Fix PYRE_MISSING_ANNOTATIONS issues in fbcode/pearl/user_envs/wrappers/sparse_reward.py

Differential Revision: D101152148

fbshipit-source-id: d67681f8008cc74a99c5f3eea9ab47fda1ac37f0

Author: generatedunixname1401739305002972

pearl/user_envs/wrappers/sparse_reward.py

@@ -6,6 +6,8 @@
 #
 # pyre-ignore-all-errors
 
+from typing import Any
+
 import numpy as np
 
 try:
@@ -69,7 +71,7 @@ def step(self, action):
         sparse_reward = self.sparse_reward()
         return obs, sparse_reward, terminated, truncated, info
 
-    def sparse_reward(self):
+    def sparse_reward(self) -> bool:
         position, velocity = self.env.state
         return bool(
             position >= self.env.goal_position and velocity >= self.env.goal_velocity
@@ -82,7 +84,7 @@ class PuckWorldSparseRewardWrapper(gym.Wrapper):
     def __init__(self, env):
         super().__init__(env)
 
-    def step(self, action):
+    def step(self, action: Any) -> tuple[Any, float, bool, bool, dict[str, Any]]:
         obs, reward, terminated, truncated, info = self.env.step(action)
         sparse_reward = self.sparse_reward()
         return obs, sparse_reward, terminated, truncated, info