Add maskable GraphPPO based on sb3_contrib.MaskablePPO + GNN for domains with graph observations #444
+3,451
−172
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
nhuet
force-pushed
the
gnn-sb3-maskable
branch
2 times, most recently
from
f52083e to
9550a07
Compare
- Use it in rollout to make them be aware of current action mask. - Add a `get_action_mask()` method to domains by default converting applicable actions space into a 0-1 numpy array, provided that the action space of each agent is an EnumerableSpace.
- inherits from Maskable - do not require anymore FullObservable from the domain to use action masking, as get_action_mask() can be called without the solver knowing about the current state (and since in rollout, the actual domain is now used) - decide whether using action masking directly in __init__() so that using_applicable_actions() can be overriden properly - use common functions for unwrap_obs and wrap_action in solver and wrapper environment to avoid code duplication - use domain.get_action_mask() to convert applicable actions into a mask (the method is more efficient as not calling get_applicable_actions() for each actions)
This is more memory sufficient for only 0-1's. And seems to be the standard for action mask at least for ray.rllib, as shown in `action_mask_key` documentation at https://docs.ray.io/en/latest/rllib/rllib-training.html
- we reuse our stable_baselines3 wrapper
- the policy is extracting features from the graph with a GNN
- the GNN is using pytorch-geometric
- We subclass
- ActorCriticPolicy:
- feature extractor = gnn
- custom conversion of observation to torch to convert into
torch_geometric.data.Data
- PPO to handle properly
- observation conversion
- rollout buffer
- Current limitations:
- we extract a fixed number of features (independent of edge/node
numbers) for now as we end with a feature reduction layer connected
to a classic mlp (not knowning anything about the current graph structure)
- User input: the user can define (and default choices are made else)
- the gnn (default to a 2 layers GCN), taking as inputs w.r.t torch_geometric conventions:
- x: nodes features
- edge_index: edge indices or sparse transposed adjency matrix
- edge_attr (optional): edges features
- edge_weight (optional): edge weights (taken from first dimension
of edge_attr)
- the feature reduction layer from the gnn output to the fixed number of features
(default to global_max_pool + linear layer + relu)
We also introduce a multiinput policy to take into account static graph
features. The observation space is a DictSpace whose subspaces can
contain some Graph spaces.
nhuet
force-pushed
the
gnn-sb3-maskable
branch
from
9550a07 to
b126b46
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
MaskableGraphPPOfromMaskablePPO(and alsoGraphOnPolicyAlgorithmto reuse work done forGraphPPO)MaskableGNNActorCriticPolicyfromMaskableActorCriticPolicy(and similarly from_BaseGNNActorCriticPolicyto share code withGNNActorCriticPolicy), idem forMaskableMultiInputGNNActorCriticPolicy.use_action_maskingin order to wrap the domain in an environment exposingaction_masksmethod (redirecting towardsdomain.get_action_mask())GraphMazeandGraphJspDomainso that the solver propose only applicable actions