Fix typing issue in python 3.9

HumanCompatibleAI · Sep 12, 2023 · d5a0a46 · d5a0a46
1 parent a7954c2
commit d5a0a46
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/seals/util.py b/src/seals/util.py
@@ -138,8 +138,8 @@ class BoxRegion:
 
 
 class MaskScoreWrapper(
-    gym.Wrapper[npt.NDArray, ActType, npt.NDArray, ActType],
-    Generic[ActType],
+    gym.Wrapper[ObsType, ActType, npt.NDArray, ActType],
+    Generic[ObsType, ActType],
 ):
     """Mask a list of box-shaped regions in the observation to hide reward info.
 
@@ -150,7 +150,7 @@ class MaskScoreWrapper(
 
     def __init__(
         self,
-        env: gym.Env,
+        env: gym.Env[ObsType, ActType],
         score_regions: MaskedRegionSpecifier,
         fill_value: Union[float, Sequence[float]] = 0,
     ):