Support for VecMonitor for gym3-style environments

[vwxyzjn]

Description

Gym3-style environments such a procgen directly produces the vectorized environments, so there is no chance of passing a Monitor wrapper to the envs creation process. However, there still is a need to record the episodic stats. One potential way to do this is through the VecMonitor wrapper as done in openai/baselines. See #302 for further related discussions.

Motivation and Context

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

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)
• I have checked that the documentation builds using make doc (required)

Was having some troubles with the make type. It seems the error was not caused by my changes, maybe?

[50/91] check stable_baselines3.common.base_class
FAILED: /home/costa/Documents/work/go/src/github.com/vwxyzjn/stable-baselines3/.pytype/pyi/stable_baselines3/common/base_class.pyi 
/home/costa/anaconda3/bin/python -m pytype.single --imports_info /home/costa/Documents/work/go/src/github.com/vwxyzjn/stable-baselines3/.pytype/imports/stable_baselines3.common.base_class.imports --module-name stable_baselines3.common.base_class -V 3.8 -o /home/costa/Documents/work/go/src/github.com/vwxyzjn/stable-baselines3/.pytype/pyi/stable_baselines3/common/base_class.pyi --analyze-annotated --nofail --quick /home/costa/Documents/work/go/src/github.com/vwxyzjn/stable-baselines3/stable_baselines3/common/base_class.py
File "/home/costa/Documents/work/go/src/github.com/vwxyzjn/stable-baselines3/stable_baselines3/common/base_class.py", line 627, in load: Can't instantiate BaseAlgorithm with abstract methods _setup_model, learn [not-instantiable]
File "/home/costa/Documents/work/go/src/github.com/vwxyzjn/stable-baselines3/stable_baselines3/common/base_class.py", line 627, in load: Invalid keyword argument _init_setup_model to function BaseAlgorithm.__init__ [wrong-keyword-args]
         Expected: (self, policy, env, policy_base, learning_rate, policy_kwargs, tensorboard_log, verbose, device, support_multi_env, create_eval_env, monitor_wrapper, seed, use_sde, sde_sample_freq, supported_action_spaces)
  Actually passed: (self, policy, env, device, _init_setup_model)

For more details, see https://google.github.io/pytype/errors.html.
ninja: build stopped: subcommand failed.
Leaving directory '/home/costa/Documents/work/go/src/github.com/vwxyzjn/stable-baselines3/.pytype'
make: *** [Makefile:8: type] Error 1

Given this new feature, the following two ways of initializing vectorized envs should be equivalent, except the latter is much faster when used with higher number of parallel envs.

from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv
import gym
def make_env(gym_id):
    def thunk():
        env = gym.make(gym_id, render_mode='rgb_array')
        env = Monitor(env)
        return env
    return thunk
envs = DummyVecEnv([make_env('procgen-starpilot-v0')])

from procgen import ProcgenEnv
from stable_baselines3.common.vec_env import VecExtractDictObs, VecMonitor
venv = ProcgenEnv(num_envs=1, env_name='starpilot')
venv = VecExtractDictObs(venv, "rgb")
venv = VecMonitor(venv=venv)

[vwxyzjn]

Just saw the CI failed, however the doc build successfully on my end. Hmmm.

(base) ~/.../vwxyzjn/stable-baselines3 >>> cd docs                                                                                                                                             
(base) ~/.../stable-baselines3/docs >>> make html                                                                                                                                              
Running Sphinx v3.1.2

** (process:21786): WARNING **: 19:33:00.337: Error loading plugin: libhspell.so.0: cannot open shared object file: No such file or directory


** (process:21786): WARNING **: 19:33:00.337: Error loading plugin: libnuspell.so.4: cannot open shared object file: No such file or directory


** (process:21786): WARNING **: 19:33:00.337: Error loading plugin: libvoikko.so.1: cannot open shared object file: No such file or directory


** (process:21786): WARNING **: 19:33:00.337: Error loading plugin: libaspell.so.15: cannot open shared object file: No such file or directory

Initializing Spelling Checker 7.1.0
loading pickled environment... done
building [mo]: targets for 0 po files that are out of date
building [html]: targets for 0 source files that are out of date
updating environment: 0 added, 0 changed, 0 removed
looking for now-outdated files... none found
no targets are out of date.
build succeeded.

The HTML pages are in _build/html.

[araffin]

The error for the doc:

/home/runner/work/stable-baselines3/stable-baselines3/stable_baselines3/common/vec_env/vec_monitor.py:docstring of stable_baselines3.common.vec_env.vec_monitor.VecMonitor:14:Unexpected indentation.

make sure to have the latest package version of sphinx ;)

[araffin]

you should probably check at the very beginning if a monitor wrapper is already present or not (cf what we do with evaluate_policy) otherwise it would overwrite it.

[vwxyzjn]

I am not sure how to do the check exactly...

[araffin]

you may take a look at https://github.com/DLR-RM/stable-baselines3/blob/master/tests/test_vec_envs.py#L432
and https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/evaluation.py#L60

[araffin]

Overall look good, please add some tests and the csv logging feature ;)

Note: we are using rst for the dosctring format, not markdown unfortunately

[vwxyzjn]

@araffin Thank you for the very helpful review. I have addressed most concerns raised in your comments. Please take a look when you have a chance :)

[araffin]

as a general remark, as mentioned in the PR template and contributing guide, please update the changelog accordingly too

[vwxyzjn]

Hey, hope you all are doing well. Just checking is there anything else you would like me to add for this PR.

[araffin]

Hey, hope you all are doing well. Just checking is there anything else you would like me to add for this PR.

yep, I think there are still some comments not addressed...

EDIT: apparently, there is one left and then I may give a try to simplify the imports

[vwxyzjn]

apparently, there is one left and then I may give a try to simplify the imports

Sounds good :) Let me know if I can provide any help.

[benblack769]

I just got another feature request for this PR here: Farama-Foundation/PettingZoo#363.

The background is that PettingZoo has support for stable baselines by converting a pettingzoo multiagent environment into an equivalent vector environment https://github.com/PettingZoo-Team/SuperSuit#parallel-environment-vectorization. So this PR will enable reward logging for multiagent environments with SB3, which is clearly very valuable.

[vwxyzjn]

Just incorporated the fix from #379

[vwxyzjn]

add a warning in the VecEnv doc

I was gonna add a warning in the VecEnv doc but wasn't sure how to add it. Especially because the gym3 vec envs is kind of a separate entities from VecEnv...

[araffin]

I've updated the PR and wanted to push the commits on my side, but I need write access to https://github.com/vwxyzjn/stable-baselines3 ...

[araffin]

LGTM, thanks =)

[vwxyzjn]

Thanks!