Fix the wrapper type hints

[pseudo-rnd-thoughts]

As noted in DLR-RM/stable-baselines3#1327 and tianshou, the wrapper type hints were raising type checker issues.
This PR solves those issues by changing the wrapper type hints to
class Wrapper(Env[WrapperObsType, WrapperActType], Generic[WrapperObsType, WrapperActType, ObsType, ActType]):

For example type hinting that passes mypy and pyright type checkers

# pyright: strict

from typing import TypeVar, Generic
from abc import abstractmethod

ObsType = TypeVar("ObsType")
ActType = TypeVar("ActType")

class Env(Generic[ObsType, ActType]):


    @abstractmethod
    def step(self, action: ActType) -> ObsType:
        pass

    @abstractmethod
    def reset(self) -> ObsType:
        pass


WrapperObsType = TypeVar("WrapperObsType")
WrapperActType = TypeVar("WrapperActType")
WrappedEnvObsType = TypeVar("WrappedEnvObsType")
WrappedEnvActType = TypeVar("WrappedEnvActType")


class Wrapper(Env[WrapperObsType, WrapperActType], Generic[WrapperObsType, WrapperActType, WrappedEnvObsType, WrappedEnvActType]):

    def __init__(self, env: Env[WrappedEnvObsType, WrappedEnvActType]) -> None:
        self.env = env

    @abstractmethod
    def step(self, action: WrapperActType) -> WrapperObsType:
        pass

    @abstractmethod
    def reset(self) -> WrapperObsType:
        pass


class ObservationWrapper(Wrapper[WrapperObsType, WrappedEnvActType, WrappedEnvObsType, WrappedEnvActType]):
    
    def __init__(self, env: Env[WrappedEnvObsType, WrappedEnvActType]) -> None:
        super().__init__(env)

    def step(self, action: WrappedEnvActType) -> WrapperObsType:
        return self.transform_observation(self.env.step(action))

    @abstractmethod
    def transform_observation(self, observation: WrappedEnvObsType) -> WrapperObsType:
        pass


class ActionWrapper(Wrapper[WrappedEnvObsType, WrapperActType, WrappedEnvObsType, WrappedEnvActType]):

    def __init__(self, env: Env[WrappedEnvObsType, WrappedEnvActType]) -> None:
        super().__init__(env)

    def step(self, action: WrapperActType) -> WrappedEnvObsType:
        return super().step(action)
    
    @abstractmethod
    def transform_action(self, action: WrapperActType) -> WrappedEnvActType:
        pass


class RewardWrapper(Wrapper[WrappedEnvObsType, WrappedEnvActType, WrappedEnvObsType, WrappedEnvActType]):

    def __init__(self, env: Env[WrappedEnvObsType, WrappedEnvActType]) -> None:
        super().__init__(env)

    def step(self, action: WrappedEnvActType) -> WrappedEnvObsType:
        return super().step(action)

    @abstractmethod
    def transform_reward(self, reward: float):
        pass

class ObsStringToIntWrapper(ObservationWrapper[int, ActType, str]):

    def __init__(self, env: Env[str, ActType]) -> None:
        super().__init__(env)

    def transform_observation(self, observation: str) -> int:
        return int(observation)


class ActionStringToIntWrapper(ActionWrapper[ObsType, str, int]):

    def __init__(self, env: Env[ObsType, int]) -> None:
        super().__init__(env)

    def transform_action(self, action: str) -> int:
        return int(action)

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated
gymnasium	✅ Ready (Inspect)	Visit Preview	💬 Add your feedback	Feb 18, 2023 at 8:39PM (UTC)

[pseudo-rnd-thoughts]

@araffin Has found that this PR fixes the type hinting issues discovered in SB3, therefore, merging