Fix RandomNumberGenerator pickling

[jjyyxx]

Reason

gym 0.22.0 defines a RandomNumberGenerator class inheriting np.random.Generator.
np.random.Generator defines __reduce__, but it's hard-coded to return a Generator instead of its subclass RandomNumberGenerator. Thus, sampling from a Space will be broken after pickling and unpickling, due to using the deprecated methods defined in RandomNumberGenerator. (Which further breaks some RL frameworks like RLlib relying on pickling for distributed training)

Example

Before

>>> import gym, pickle
>>> space = gym.spaces.Discrete(500)
>>> rng = pickle.loads(pickle.dumps(space.np_random))
>>> rng
Generator(PCG64) at _
>>> rng.randint(2)
AttributeError: 'numpy.random._generator.Generator' object has no attribute 'randint'

After

>>> import gym, pickle
>>> space = gym.spaces.Discrete(500)
>>> rng = pickle.loads(pickle.dumps(space.np_random))
>>> rng
RandomNumberGenerator(PCG64) at _
>>> rng.randint(2)
1

[RedTachyon]

Can you add some tests to check that this will indeed work?

[jjyyxx]

Test added 342a0a2. Thanks for pointing it out.

[jkterry1]

@jjyyxx can you please fix tests?

[jjyyxx]

I pushed a fix.

[jkterry1]

@jjyyxx tests still aren't passing

[jjyyxx]

My bad. Didn't set up pre-commit check. Now I hope everything is ok.

[RedTachyon]

I don't think this actually solves the issue though? Please add a test that pickles an actual space (a Discrete and a Box - at least two tests) and then loads it in the way that causes the bug reported here. So:

1. Initialize space
2. Sample from space
3. Pickle space (to a file and to a string - need to cover both cases with the test)
4. Unpickle it as space2
5. Sample from space2

With the current code, I'm still getting the same exception

[jjyyxx]

@RedTachyon Thanks for your feedback. Based on your suggestion, I write a test case like this:

from gym.spaces import Box, Discrete
import tempfile
@pytest.mark.parametrize("space", [
    Discrete(100),
    Box(low=0.0, high=1.0, shape=(2, 2)),
])
def test_space_pickle(space):
    space.sample()  # To initialize the RNG, or access space.np_random

    # Pickle and unpickle with a string
    pickled = pickle.dumps(space)
    space2 = pickle.loads(pickled)

    # Pickle and unpickle with a file
    with tempfile.TemporaryFile() as f:
        pickle.dump(space, f)
        f.seek(0)
        space3 = pickle.load(f)

    value = space.sample()
    value2 = space2.sample()
    value3 = space3.sample()
    assert np.array_equal(value, value2)
    assert np.array_equal(value, value3)

And my fix could pass this test. But as a first-time contributor, I'm not confident where to put this test, e.g.

1. add into test_spaces.py test_roundtripping function
2. a standalone function in test_spaces.py
3. elsewhere

So I did not commit this test case yet.

Besides, I'm not quite sure how to reproduce the last sentence you mentioned:

With the current code, I'm still getting the same exception

The only pitfall I found was you must sample or access space.np_random explicitly before pickle to ensure random seed reproducibility (does not affect pickling), due to _np_random lazy initialization. But I think it may be related to the space API, and out of scope for this RNG fix.

[jkterry1]

Please put it in the second location

[RedTachyon]

I might have messed something up with my test, because now I'm rerunning the same code and works.

My main concern was for the space pickle-unpickle cycle to not crash, and this seems to work now, so imo this can be merged and then we'll see if SB3 is happy

[jjyyxx]

I met this issue when working with RLlib. In my test, it works after applying this fix.

[carlosluis]

I might have messed something up with my test, because now I'm rerunning the same code and works.

My main concern was for the space pickle-unpickle cycle to not crash, and this seems to work now, so imo this can be merged and then we'll see if SB3 is happy

I ran the previously failing SB3 tests again and now they pass, so this seems to fix it.