Fix the terminal observation for gym3-style environments

[vwxyzjn]

Description

This PR fixes the terminal observation for gym3-style environments. Combined with the #311, SB3 would support training agents in procgen.

import gym
from procgen import ProcgenEnv
from stable_baselines3 import A2C
from stable_baselines3.common.vec_env import VecEnvWrapper, VecNormalize
import time
import numpy as np

class VecExtractDictObs(VecEnvWrapper):
    def __init__(self, venv, key):
        self.key = key
        super().__init__(venv=venv,
            observation_space=venv.observation_space.spaces[self.key])

    def reset(self):
        obs = self.venv.reset()
        return obs[self.key]

    def step_wait(self):
        obs, reward, done, info = self.venv.step_wait()
        return obs[self.key], reward, done, info

class VecMonitor(VecEnvWrapper):
    def __init__(self, venv):
        VecEnvWrapper.__init__(self, venv)
        self.eprets = None
        self.eplens = None
        self.epcount = 0
        self.tstart = time.time()

    def reset(self):
        obs = self.venv.reset()
        self.eprets = np.zeros(self.num_envs, 'f')
        self.eplens = np.zeros(self.num_envs, 'i')
        return obs

    def step_wait(self):
        obs, rews, dones, infos = self.venv.step_wait()
        self.eprets += rews
        self.eplens += 1

        newinfos = list(infos[:])
        for i in range(len(dones)):
            if dones[i]:
                info = infos[i].copy()
                ret = self.eprets[i]
                eplen = self.eplens[i]
                epinfo = {'r': ret, 'l': eplen, 't': round(time.time() - self.tstart, 6)}
                info['episode'] = epinfo
                self.epcount += 1
                self.eprets[i] = 0
                self.eplens[i] = 0
                newinfos[i] = info
        return obs, rews, dones, newinfos


venv = ProcgenEnv(num_envs=10, env_name='starpilot')
venv = VecExtractDictObs(venv, "rgb")
env = VecMonitor(venv=venv)
model = A2C('CnnPolicy', env, verbose=1, device="cpu")
model.learn(total_timesteps=1000000)

See

There are only two lines of code so I kind of just skips the checklist if that's ok...

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)

Note: You can run most of the checks using make commit-checks.

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

[araffin]

This PR fixes the terminal observation for gym3-style environments. Combined with the #311, SB3 would support training agents in procgen.
There are only two lines of code so I kind of just skips the checklist if that's ok...

I still have mixed feelings about that, as it would introduce wrong terminal observations silently... (I still think this should be also fixed upstream in gym3).

As it is only two lines of code maybe you could merge it with #311 no? (but still add it to the changelog and add a warning in the VecEnv doc)