Add EnvPool support (as VecEnv subclass)

[qgallouedec]

Description

Based on #307, but use VecEnv subclass instead of VecEnvWrapper.
VecEnvWrapper assumes the wrapped env of an instance of VecEnv, which is not the case of EnvPool. IMO, using the VecEnv class directly is therefore a better option, and the implementation is not more complex.

Motivation and Context

• I have raised an issue to propose this change (required for new features and bug fixes)

closes #241

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)

Checklist:

• I've read the CONTRIBUTION guide (required)
• I have updated the changelog accordingly (required).
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have reformatted the code using make format (required)
• I have checked the codestyle using make check-codestyle and make lint (required)
• I have ensured make pytest and make type both pass. (required)

Note: we are using a maximum length of 127 characters per line

[qgallouedec]

Here are some speed results with Breakout and PPO. I can't understand why envpool is only slightly faster than dummy vec env. @araffin, @Trinkle23897 any insight?

(y axis is fps)

[qgallouedec]

@araffin what do you think about moving this to sb3 (insert it to make_vec_env)?
And make envpool an optional dependency of sb3.

[araffin]

what do you think about moving this to sb3 (insert it to make_vec_env)?

not sure about it.
My main concern is that it only supports half the thing that a VecEnv can do (notably, it doesn't support gym wrappers).
I think it would make more sense in sb3 contrib or here in the RL Zoo (especially as it is a package now, so people can re-use it easily if needed).

[qgallouedec]

My main concern is that it only supports half the thing that a VecEnv can do

Makes sense, we will need to update the sb3 doc btw.

[qgallouedec]

@araffin, before merging, we need to explain this #355 (comment)

[araffin]

I don't think we have vec_env_kwargs available in the cli, but we should add it.

[araffin]

Suggested change

	if (
	"Neck" in self.env_name.gym_id
	or self.is_robotics_env(self.env_name.gym_id)
	or "parking-v0" in self.env_name.gym_id
	and len(self.monitor_kwargs) == 0 # do not overwrite custom kwargs
	):
	if (
	self.is_robotics_env(self.env_name.gym_id)
	or "parking-v0" in self.env_name.gym_id
	and len(self.monitor_kwargs) == 0 # do not overwrite custom kwargs
	):

"Neck" in self.env_name.gym_id was for personal use and should be removed.

[araffin]

Missing: a section in the readme, I will try to create a real doc soon (the readme is way too long now).

[araffin]

doc is there

[araffin]

we don't have that type hint defined in SB3?

[qgallouedec]

Yes, in base_with_env. I don't remember why I redefined it here. I probably should have imported it

[araffin]

Here are some speed results with Breakout and PPO. I can't understand why envpool is only slightly faster than dummy vec env. @araffin, @Trinkle23897 any insight?

So, it seems to depends on the env, but yes it is surprising.

I also have envpool slower than SubProcVecEnv on my machine with python 3.7 (but not as slow as reported here).

But on MuJoCo env, envpool is the fastest for me:

Envpool: ~4100 fps

CUDA_VISIBLE_DEVICES= OMP_NUM_THREADS=2 python -m rl_zoo3.train --algo a2c --env HalfCheetah-v3 --verbose 0 --eval-freq -1 -P -envpool -params n_envs:4

vs
Subproc: ~3400 fps

CUDA_VISIBLE_DEVICES= OMP_NUM_THREADS=2 python -m rl_zoo3.train --algo a2c --env HalfCheetah-v3 --verbose 0 --eval-freq -1 -P -params n_envs:4 --vec-env subproc

vs

Dummy: ~3600 fps

CUDA_VISIBLE_DEVICES= OMP_NUM_THREADS=2 python -m rl_zoo3.train --algo a2c --env HalfCheetah-v3 --verbose 0 --eval-freq -1 -P -params n_envs:4 --vec-env dummy

[Trinkle23897]

Can you use py-spy to have a profiling result?

[araffin]

Can you use py-spy to have a profiling result?

Profiles:

This is for sub process:

For envpool, I think py-spy had a segfault:

[Trinkle23897]

It seems like there are a lot of "importlib" calls in rl_zoo3.train with envpool. That should not be the case and I believe something is wrong. Did you accidentally call envpool with multiple python processes/threads?