Is your feature request related to a problem? Please describe
Currently, wrappers don't inherit type hinting from the environments they wrap. For example, consider the return types of the reset method. As you can see when using the environment directly, type hinting works as expected.
However when using a wrapper around the environment, a couple problems show up. First, we are told that we need to annotate the wrapped environment. Second, type information about state and timestep.observation are lost.
game_2048 = Wrapper(Game2048()) # error: Need type annotation for "game_2048"
state, timestep = game_2048.reset(key)
reveal_type(state) # Any
reveal_type(timestep.observation) # Any
Explicitly annotating the environment fixes the error and restores the type hint on the state variable. However, there is still no type hint for timestep.observation.
This will also provide type hints for the state variable. Providing type hints for timestep.observation will probably require adding a generic Observation typevar to the Environment class.
Is your feature request related to a problem? Please describe
Currently, wrappers don't inherit type hinting from the environments they wrap. For example, consider the return types of the
reset
method. As you can see when using the environment directly, type hinting works as expected.However when using a wrapper around the environment, a couple problems show up. First, we are told that we need to annotate the wrapped environment. Second, type information about
state
andtimestep.observation
are lost.Explicitly annotating the environment fixes the error and restores the type hint on the
state
variable. However, there is still no type hint fortimestep.observation
.I believe a similar problem exists for most of the methods on the
Wrapper
class.Describe the solution you'd like
Ideally, the following code should just give proper type hints.
The first issue of needing to explicitly annotate the wrapped environment can be solved by modifying the
__init__
method as followsThis will also provide type hints for the
state
variable. Providing type hints fortimestep.observation
will probably require adding a genericObservation
typevar to theEnvironment
class.