lit_sb3.wrappers.action_masker
[docs]
module
lit_sb3.wrappers.action_masker
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33 | from collections.abc import Callable
from typing import Union
import gymnasium as gym
import numpy as np
class ActionMasker(gym.Wrapper):
"""Env wrapper providing the method required to support masking.
Exposes a method called action_masks(), which returns masks for the wrapped env.
This wrapper is not needed if the env exposes the expected method itself.
Args:
env: the Gym environment to wrap
action_mask_fn: A function that takes a Gym environment and returns an action mask,
or the name of such a method provided by the environment.
"""
def __init__(self, env: gym.Env, action_mask_fn: str | Callable[[gym.Env], np.ndarray]) -> None:
super().__init__(env)
if isinstance(action_mask_fn, str):
found_method = getattr(self.env, action_mask_fn)
if not callable(found_method):
raise TypeError(f"Environment attribute {action_mask_fn} is not a method")
self._action_mask_fn = found_method
else:
self._action_mask_fn = action_mask_fn
def action_masks(self) -> np.ndarray:
return self._action_mask_fn(self.env)
|