PyTorch 1.9 Compatibility: IterableDataset Type Annotations

[lebrice]

Creating this issue as a follow-up to the discussion with @ejguan on the closed PR #233 :

PyTorch 1.9 checks the type annotations on the IterableDataset subclasses. This is currently causing issues with the annotations of the various subclasses of Sequoia's EnvDataset class, which inherits from IterableDataset.

This EnvDataset class can be found here, but here is a simple version of what it is meant to be like:

Note: The actual EnvDataset class is a bit different from this, and doesn't use a generator, but this is just to illustrate the kind of interaction we have with IterableDataset:

from torch.utils.data import IterableDataset
import gym
import numpy as np
from typing import Generator, Optional, Generic, TypeVar, Tuple, Union, Iterator

ObservationType = TypeVar("ObservationType")
ActionType = TypeVar("ActionType")
RewardType = TypeVar("RewardType")


class EnvDataset(
    gym.Wrapper,
    IterableDataset,
):
    def __init__(self, env: gym.Env):
        super().__init__(env=env)
        self.env = env
        self._iterator: Optional[Generator] = None

    def __iter__(self):
        if self._iterator:
            self._iterator.close()
        self._iterator = self.iterate(self.env)
        return self._iterator

    @staticmethod
    def iterate(env) -> Generator[Union[ObservationType, RewardType], ActionType, None]:
        """Iterator / generator for a gym.Env."""
        try:
            observations = env.reset()
            done = False
            while not done:
                actions = yield observations
                if actions is None:
                    raise RuntimeError("Need to send an action after each observation.")
                observations, rewards, done, info = env.step(actions)
                yield rewards
        except GeneratorExit:
            print("closing")

    def send(self, actions: ActionType) -> RewardType:
        return self._iterator.send(actions)

env = gym.make("CartPole-v0")
env_dataset = EnvDataset(env)

for i, obs in enumerate(env_dataset):
    action = env_dataset.action_space.sample()
    rewards = env_dataset.send(action)
    print(i, obs, action)

[ejguan]

Just tested the EnvDataset class itself and it should not break anything. Please correct me if I am wrong, I didn't find any subclass of EnvDataset in your repo. D
BTW, can you share the python version?
Edit: I did try python3.7 and python3.9 without any error.

[lebrice]

Oh my mistake, I meant subclasses of IterableWrapper, a type of gym.Wrapper used around these EnvDatasets.

Indeed this example doesn't have any errors! I'm pretty sure this issue isn't that hard to fix, I probably need to make sure all the __iter__ methods don't have any type annotation or just have an Iterator[Any] annotation so the errors go away.
Python 3.8 at the moment, but Sequoia is meant to work for 3.7+.

[ejguan]

I see. I have a super simple fix for you, remove the bound for EnvType here:

Sequoia/sequoia/common/gym_wrappers/utils.py

Line 220 in 517ccd2

EnvType = TypeVar("EnvType", bound=gym.Env)

It's true that it is too hard to figure out without context.

For example of MeasureSLPerformanceWrapper:

Sequoia/sequoia/settings/sl/wrappers/measure_performance.py

Lines 32 to 36 in fc767ee

	class MeasureSLPerformanceWrapper(
	MeasurePerformanceWrapper,
	# MeasurePerformanceWrapper[PassiveEnvironment] # Python 3.7
	# MeasurePerformanceWrapper[PassiveEnvironment, ClassificationMetrics] # Python 3.8+
	):

It inherits MeasurePerformanceWrapper . And, you can find you have specified the type is EnvType with bound of gym.env as shown below.

Sequoia/sequoia/common/gym_wrappers/measure_performance.py

Lines 13 to 14 in fc767ee

	class MeasurePerformanceWrapper(
	IterableWrapper[EnvType], Generic[EnvType, MetricsType], ABC

We just introduced a hard check for the return type of __iter__ in 1.9. But, you just specified the return type as below, which is not a subtype of your EnvType IMO. So, the Error raises.

Sequoia/sequoia/settings/sl/wrappers/measure_performance.py

Line 144 in fc767ee

def __iter__(self) -> Iterator[Tuple[Observations, Optional[Rewards]]]:

But, I agree this check may be too strict. I may try to release the check for iterator in the next release. Thank you for your feedback!

BTW, during my testing, I did find some interesting part like:

Sequoia/sequoia/settings/rl/continual/environment.py

Lines 80 to 82 in 84c35e6

	class GymDataLoader(
	ActiveEnvironment[ObservationType, ActionType, RewardType], IterableWrapper, Iterable
	):

The GymDataLoader is inheriting both DataLoader and IterableDataset, which is kind unexpected even though it doesn't break anything.

[ejguan]

Hey @lebrice ,

Since it's been a while, just want to follow up if there is any other concern about it?

[lebrice]

@ejguan Hey! No its all good, I just havent taken the time to do it. I'll probably get this sorted out soon though.

Thanks again!