from typing import Dict, Optional, Tuple, Union import gymnasium as gym from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule from ray.rllib.core.columns import Columns from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig from ray.rllib.core.rl_module.rl_module import RLModule from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_torch from ray.rllib.utils.torch_utils import FLOAT_MIN from ray.rllib.utils.typing import TensorType torch, nn = try_import_torch() class ActionMaskingRLModule(RLModule): """An RLModule that implements an action masking for safe RL. This RLModule implements action masking to avoid unsafe/unwanted actions dependent on the current state (observations). It does so by using an environment generated action mask defining which actions are allowed and which should be avoided. The action mask is extracted from the environment's `gymnasium.spaces.dict.Dict` observation and applied after the module's `forward`-pass to the action logits. The resulting action logits prevent unsafe/unwanted actions to be sampled from the corresponding action distribution. Note, this RLModule is implemented for the `PPO` algorithm only. It is not guaranteed to work with other algorithms. Furthermore, not that for this module to work it requires an environment with a `gymnasium.spaces.dict.Dict` observation space containing tow key, `"action_mask"` and `"observations"`. """ @override(RLModule) def __init__( self, *, observation_space: Optional[gym.Space] = None, action_space: Optional[gym.Space] = None, inference_only: Optional[bool] = None, learner_only: bool = False, model_config: Optional[Union[dict, DefaultModelConfig]] = None, catalog_class=None, **kwargs, ): # If observation space is not of type `Dict` raise an error. if not isinstance(observation_space, gym.spaces.dict.Dict): raise ValueError( "This RLModule requires the environment to provide a " "`gym.spaces.Dict` observation space of the form: \n" " {'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,))," " 'observation_space': self.observation_space}" ) # While the environment holds an observation space that contains, both, # the action mask and the original observation space, the 'RLModule' # receives only the `"observation"` element of the space, but not the # action mask. self.observation_space_with_mask = observation_space self.observation_space = observation_space["observations"] # Keeps track if observation specs have been checked already. self._checked_observations = False # The DefaultPPORLModule, in its constructor will build networks for the # original observation space (i.e. without the action mask). super().__init__( observation_space=self.observation_space, action_space=action_space, inference_only=inference_only, learner_only=learner_only, model_config=model_config, catalog_class=catalog_class, **kwargs, ) class ActionMaskingTorchRLModule(ActionMaskingRLModule, PPOTorchRLModule): @override(PPOTorchRLModule) def setup(self): super().setup() # We need to reset here the observation space such that the # super`s (`PPOTorchRLModule`) observation space is the # original space (i.e. without the action mask) and `self`'s # observation space contains the action mask. self.observation_space = self.observation_space_with_mask @override(PPOTorchRLModule) def _forward_inference( self, batch: Dict[str, TensorType], **kwargs ) -> Dict[str, TensorType]: # Preprocess the original batch to extract the action mask. action_mask, batch = self._preprocess_batch(batch) # Run the forward pass. outs = super()._forward_inference(batch, **kwargs) # Mask the action logits and return. return self._mask_action_logits(outs, action_mask) @override(PPOTorchRLModule) def _forward_exploration( self, batch: Dict[str, TensorType], **kwargs ) -> Dict[str, TensorType]: # Preprocess the original batch to extract the action mask. action_mask, batch = self._preprocess_batch(batch) # Run the forward pass. outs = super()._forward_exploration(batch, **kwargs) # Mask the action logits and return. return self._mask_action_logits(outs, action_mask) @override(PPOTorchRLModule) def _forward_train( self, batch: Dict[str, TensorType], **kwargs ) -> Dict[str, TensorType]: # Run the forward pass. outs = super()._forward_train(batch, **kwargs) # Mask the action logits and return. return self._mask_action_logits(outs, batch["action_mask"]) @override(ValueFunctionAPI) def compute_values(self, batch: Dict[str, TensorType], embeddings=None): # Check, if the observations are still in `dict` form. if isinstance(batch[Columns.OBS], dict): # Preprocess the batch to extract the `observations` to `Columns.OBS`. action_mask, batch = self._preprocess_batch(batch) # NOTE: Because we manipulate the batch we need to add the `action_mask` # to the batch to access them in `_forward_train`. batch["action_mask"] = action_mask # Call the super's method to compute values for GAE. return super().compute_values(batch, embeddings) def _preprocess_batch( self, batch: Dict[str, TensorType], **kwargs ) -> Tuple[TensorType, Dict[str, TensorType]]: """Extracts observations and action mask from the batch Args: batch: A dictionary containing tensors (at least `Columns.OBS`) Returns: A tuple with the action mask tensor and the modified batch containing the original observations. """ # Check observation specs for action mask and observation keys. self._check_batch(batch) # Extract the available actions tensor from the observation. action_mask = batch[Columns.OBS].pop("action_mask") # Modify the batch for the `DefaultPPORLModule`'s `forward` method, i.e. # pass only `"obs"` into the `forward` method. batch[Columns.OBS] = batch[Columns.OBS].pop("observations") # Return the extracted action mask and the modified batch. return action_mask, batch def _mask_action_logits( self, batch: Dict[str, TensorType], action_mask: TensorType ) -> Dict[str, TensorType]: """Masks the action logits for the output of `forward` methods Args: batch: A dictionary containing tensors (at least action logits). action_mask: A tensor containing the action mask for the current observations. Returns: A modified batch with masked action logits for the action distribution inputs. """ # Convert action mask into an `[0.0][-inf]`-type mask. inf_mask = torch.clamp(torch.log(action_mask), min=FLOAT_MIN) # Mask the logits. batch[Columns.ACTION_DIST_INPUTS] += inf_mask # Return the batch with the masked action logits. return batch def _check_batch(self, batch: Dict[str, TensorType]) -> Optional[ValueError]: """Assert that the batch includes action mask and observations. Args: batch: A dicitonary containing tensors (at least `Columns.OBS`) to be checked. Raises: `ValueError` if the column `Columns.OBS` does not contain observations and action mask. """ if not self._checked_observations: if "action_mask" not in batch[Columns.OBS]: raise ValueError( "No action mask found in observation. This `RLModule` requires " "the environment to provide observations that include an " "action mask (i.e. an observation space of the Dict space " "type that looks as follows: \n" "{'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,))," "'observations': self.observation_space}" ) if "observations" not in batch[Columns.OBS]: raise ValueError( "No observations found in observation. This 'RLModule` requires " "the environment to provide observations that include the original " "observations under a key `'observations'` in a dict (i.e. an " "observation space of the Dict space type that looks as follows: \n" "{'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,))," "'observations': }" ) self._checked_observations = True