[Feature Request] VecMaskWrapper for MaskablePPO

[CAI23sbP]

🚀 Feature

From the issue.

@araffin !
I am planning to use a Masknet (a custom neural network for masking invalid actions) that requires batch processing

Here is my test code. Environment name is 'CartPole-v1'

My library version is 'sb3_contrib==2.1.0 , stable-baselines==2.1.0'

1. modify a original code

def get_action_masks(env: GymEnv) -> np.ndarray:
    """
    Checks whether gym env exposes a method returning invalid action masks
    :param env: the Gym environment to get masks from
    :return: A numpy array of the masks
    """

    if isinstance(env, VecEnv):
        return env.get_attr(EXPECTED_METHOD_NAME)
    else:
        return getattr(env, EXPECTED_METHOD_NAME)
    
    
def is_masking_supported(env: GymEnv) -> bool:
    """
    Checks whether gym env exposes a method returning invalid action masks

    :param env: the Gym environment to check
    :return: True if the method is found, False otherwise
    """
    if isinstance(env, VecEnv):
        try:
            env.get_attr(EXPECTED_METHOD_NAME)
            
            return True
        except AttributeError:
            return False
    else:
        return hasattr(env, EXPECTED_METHOD_NAME)

2. Add a VecMaskWrapper : masking a invalid action

from typing import List
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn, VecEnvWrapper
import numpy as np 

EXPECTED_METHOD_NAME = 'action_masks'
class VecMaskWrapper(VecEnvWrapper):
    def __init__(
        self,
        venv: VecEnv,
    ):
        VecEnvWrapper.__init__(self, venv)
        temp_env = venv
        actions = temp_env.action_space.n
        self.num_envs = temp_env.num_envs
        self.possible_actions = np.arange(actions)
        self.all_valid_mask = np.ones((self.num_envs, actions)).astype(np.bool_)

    def reset(self) -> VecEnvObs:
        self.observations = self.venv.reset()
        return self.observations

    def action_masks(self) -> List[bool]:
        """ 
            https://www.gymlibrary.dev/environments/classic_control/cart_pole/
            1. The pole angle can be observed between (-.418, .418) radians (or ±24°), 
            but the episode terminates if the pole angle is not in the range (-.2095, .2095) (or ±12°)

        """
        masks = np.ones_like(self.all_valid_mask, dtype=np.bool_)
        condition_1 = np.where(self.observations[:, 2] <= -0.05)[0] # left terminate
        condition_2 = np.where(self.observations[:, 2] >= 0.05)[0] # right terminate
        masks[condition_1, 1] = False 
        masks[condition_2, 0] = False 
        return masks
        
            
    def get_attr(self, attr_name, indices=None):
        if attr_name == EXPECTED_METHOD_NAME:
            return self.action_masks()
        else:
            return super().get_attr(attr_name, indices)

    def step_wait(self) -> VecEnvStepReturn:
        self.observations, rews, dones, infos = self.venv.step_wait()
        return self.observations, rews, dones, infos

3. Example test code

from stable_baselines3.common.env_util import make_vec_env
from sb3_contrib.ppo_mask.ppo_mask import MaskablePPO

if __name__ == "__main__":
    env_id = "CartPole-v1"
    vec_env = make_vec_env(env_id, n_envs= 2) 
    vec_env = VecMaskWrapper(vec_env)
    vec_env.reset()
    model = MaskablePPO("MlpPolicy", vec_env, verbose=1)
    model.learn(total_timesteps=25_000)

    obs = vec_env.reset()
    for _ in range(1000):
        action_masks = get_action_masks(vec_env)
        action, _states = model.predict(obs, action_masks= action_masks)
        obs, rewards, dones, info = vec_env.step(action)
        vec_env.render()

Motivation

No response

Pitch

No response

Alternatives

No response

Additional context

No response

Checklist

• I have checked that there is no similar issue in the repo
• If I'm requesting a new feature, I have proposed alternatives

[araffin]

Original issue: #68

@CAI23sbP I would simplify by having a hasattr() check in isinstance(env, VecEnv), or a is_vec_env_wrapped(env, VecMaskWrapper)

Could you submit a PR with tests and doc?