[Termination Truncation Update] Refactor multiproc code

[WillDudley]

Summary

My debugging skills and I have reached an impasse, and I believe this debug difficulty is pretty much the same thing that caused Jat some grief.

When running tests, there is a failure in env.reset() somewhere. In Jat's PR, https://github.com/Farama-Foundation/SuperSuit/pull/165, Ben suggests the issue is related to:

To get reset() to return info, you will also need to extract the infos returned here https://github.com/Farama-Foundation/SuperSuit/blob/master/supersuit/vector/multiproc_vec.py#L179

The above line referenced is the same line that's raising errors. It's a shape mismatch (expected less params than a func returns), but that's all I can infer from the traceback and Pycharm debugger.

Test tracebacks

1

test/test_vector/test_gym_vector.py:54 (test_gym_supersuit_equivalency)
def test_gym_supersuit_equivalency():
        env = gym.make("MountainCarContinuous-v0")
        num_envs = 3
        venv1 = concat_vec_envs_v1(env, num_envs)
        venv2 = gym_vec_env_v0(env, num_envs)
>       check_vec_env_equivalency(venv1, venv2)

test/test_vector/test_gym_vector.py:60: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
test/test_vector/test_gym_vector.py:38: in check_vec_env_equivalency
    obs1 = venv1.reset(seed=51)
supersuit/vector/concat_vec_env.py:46: in reset
    return self.concat_obs(_res_obs)
supersuit/vector/concat_vec_env.py:68: in concat_obs
    return concatenate(
/usr/lib/python3.9/functools.py:888: in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
venv/lib/python3.9/site-packages/gym/vector/utils/numpy_utils.py:50: in _concatenate_base
    return np.stack(items, axis=0, out=out)
<__array_function__ internals>:180: in stack
    ???
venv/lib/python3.9/site-packages/numpy/core/shape_base.py:433: in stack
    return _nx.concatenate(expanded_arrays, axis=axis, out=out)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

args = ([array([[array([-0.41151175,  0.        ], dtype=float32), {}]],
      dtype=object), array([[array([-0.4768474,  0. ... ], dtype=float32), {}]], dtype=object), array([[array([-0.5977746,  0.       ], dtype=float32), {}]], dtype=object)],)
kwargs = {'axis': 0, 'out': array([[0., 0.],
       [0., 0.],
       [0., 0.]], dtype=float32)}
relevant_args = [array([[array([-0.41151175,  0.        ], dtype=float32), {}]],
      dtype=object), array([[array([-0.4768474,  0.  ...,  0.       ], dtype=float32), {}]], dtype=object), array([[0., 0.],
       [0., 0.],
       [0., 0.]], dtype=float32)]

>   ???
E   TypeError: Cannot cast array data from dtype('O') to dtype('float32') according to the rule 'same_kind'

<__array_function__ internals>:180: TypeError

2

test/test_vector/test_gym_vector.py:62 (test_inital_state_dissimilarity)
def test_inital_state_dissimilarity():
        env = gym.make("CartPole-v1")
        venv = concat_vec_envs_v1(env, 2)
>       observations = venv.reset()

test/test_vector/test_gym_vector.py:66: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
supersuit/vector/concat_vec_env.py:46: in reset
    return self.concat_obs(_res_obs)
supersuit/vector/concat_vec_env.py:68: in concat_obs
    return concatenate(
/usr/lib/python3.9/functools.py:888: in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
venv/lib/python3.9/site-packages/gym/vector/utils/numpy_utils.py:50: in _concatenate_base
    return np.stack(items, axis=0, out=out)
<__array_function__ internals>:180: in stack
    ???
venv/lib/python3.9/site-packages/numpy/core/shape_base.py:433: in stack
    return _nx.concatenate(expanded_arrays, axis=axis, out=out)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

args = ([array([[array([ 0.01767829, -0.04421717, -0.04463257,  0.01125965], dtype=float32),
        {}]], dtype=object), array([[array([ 0.04047002,  0.04853883,  0.01476109, -0.02565673], dtype=float32),
        {}]], dtype=object)],)
kwargs = {'axis': 0, 'out': array([[0., 0., 0., 0.],
       [0., 0., 0., 0.]], dtype=float32)}
relevant_args = [array([[array([ 0.01767829, -0.04421717, -0.04463257,  0.01125965], dtype=float32),
        {}]], dtype=object), arra...65673], dtype=float32),
        {}]], dtype=object), array([[0., 0., 0., 0.],
       [0., 0., 0., 0.]], dtype=float32)]

>   ???
E   ValueError: Output array is the wrong shape

<__array_function__ internals>:180: ValueError

3

test/test_vector/test_gym_vector.py:78 (test_mutliproc_single_proc_equivalency)
def test_mutliproc_single_proc_equivalency():
        env = gym.make("CartPole-v1")
        num_envs = 3
        # uses single threaded vector environment
        venv1 = concat_vec_envs_v1(env, num_envs, num_cpus=0)
        # uses multiprocessing vector environment
        venv2 = concat_vec_envs_v1(env, num_envs, num_cpus=4)
>       check_vec_env_equivalency(venv1, venv2)

test/test_vector/test_gym_vector.py:86: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
test/test_vector/test_gym_vector.py:38: in check_vec_env_equivalency
    obs1 = venv1.reset(seed=51)
supersuit/vector/concat_vec_env.py:46: in reset
    return self.concat_obs(_res_obs)
supersuit/vector/concat_vec_env.py:68: in concat_obs
    return concatenate(
/usr/lib/python3.9/functools.py:888: in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
venv/lib/python3.9/site-packages/gym/vector/utils/numpy_utils.py:50: in _concatenate_base
    return np.stack(items, axis=0, out=out)
<__array_function__ internals>:180: in stack
    ???
venv/lib/python3.9/site-packages/numpy/core/shape_base.py:433: in stack
    return _nx.concatenate(expanded_arrays, axis=axis, out=out)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

args = ([array([[array([ 0.04424412, -0.01928616, -0.02450934, -0.02773815], dtype=float32),
        {}]], dtype=object), arr...ct), array([[array([-0.0488873 ,  0.02189106, -0.01688728,  0.04330887], dtype=float32),
        {}]], dtype=object)],)
kwargs = {'axis': 0, 'out': array([[0., 0., 0., 0.],
       [0., 0., 0., 0.],
       [0., 0., 0., 0.]], dtype=float32)}
relevant_args = [array([[array([ 0.04424412, -0.01928616, -0.02450934, -0.02773815], dtype=float32),
        {}]], dtype=object), arra...       {}]], dtype=object), array([[0., 0., 0., 0.],
       [0., 0., 0., 0.],
       [0., 0., 0., 0.]], dtype=float32)]

>   ???
E   ValueError: Output array is the wrong shape

<__array_function__ internals>:180: ValueError

4

  File "/home/will/PycharmProjects/SuperSuit/supersuit/vector/multiproc_vec.py", line 73, in async_loop
    observations = vec_env.reset(seed=data[0], options=data[2])
  File "/home/will/PycharmProjects/SuperSuit/supersuit/vector/concat_vec_env.py", line 46, in reset
    return self.concat_obs(_res_obs)
  File "/home/will/PycharmProjects/SuperSuit/supersuit/vector/concat_vec_env.py", line 68, in concat_obs
    return concatenate(
  File "/usr/lib/python3.9/functools.py", line 888, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/home/will/PycharmProjects/SuperSuit/venv/lib/python3.9/site-packages/gym/vector/utils/numpy_utils.py", line 50, in _concatenate_base
    return np.stack(items, axis=0, out=out)
  File "<__array_function__ internals>", line 180, in stack
  File "/home/will/PycharmProjects/SuperSuit/venv/lib/python3.9/site-packages/numpy/core/shape_base.py", line 433, in stack
    return _nx.concatenate(expanded_arrays, axis=axis, out=out)
  File "<__array_function__ internals>", line 180, in concatenate
ValueError: Output array is the wrong shape

test/test_vector/test_gym_vector.py:99 (test_multiproc_buffer)
def test_multiproc_buffer():
        num_envs = 2
        env = gym.make("CartPole-v1")
        env = concat_vec_envs_v1(env, num_envs, num_cpus=2)

>       obss = env.reset()

test/test_vector/test_gym_vector.py:105: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
supersuit/vector/multiproc_vec.py:183: in reset
    self._receive_info()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = ProcConcatVec(2)

    def _receive_info(self):
        all_data = []
        for cin in self.pipes:
            data = cin.recv()
            if isinstance(data, tuple):
                e, tb = data
                print(tb)
>               raise e
E               ValueError: Output array is the wrong shape

supersuit/vector/multiproc_vec.py:200: ValueError

Reproduction

Make sure you pip install PZ's master branch (pip install git+...;'), it is insetup.py`

Run pytest to see the failing tests.