DLR-RM / DLR-RM/stable-baselines3
[Feature Request] API for using custom distributions with existing policies
- Dominant language
- Python
- Stars
- 13.8k
- Forks
- 2.2k
- Avg merge
- 1h 35m
- Merged PRs (30d)
- 2
Description
### 🚀 Feature
The `Distribution` abstraction is defined [here](https://github.com/DLR-RM/stable-baselines3/blob/798b16aaf76ec051498a730a1d922e00184b4473/stable_baselines3/common/distributions.py#L15). But the way to introduce a new distribution to be used in the context of existing policy implementation is not obvious. In a few places, SB3 processes distributions in closed-formed manner with respect to enumerated implementations, e.g. [here](https://github.com/DLR-RM/stable-baselines3/blob/798b16aaf76ec051498a730a1d922e00184b4473/stable_baselines3/common/distributions.py#L670-L674).
Would be nice to have a hook for injecting new distribution (or providing different implementation for existing one).
### Motivation
Here's my example. I have PPO with `spaces.MultiDiscrete` action space but I want (need, probably) to have implementation of `MultiCategoricalDistribution` tailored to a very specific use case. Ideally, having as little changes to PPO and/or existing `ActorCriticPolicy` implementation as possible. The [code](https://github.com/kachayev/gym-microrts-paper-sb3/blob/48bcc7f609bb85eeffc5ae0dee5e64125abc59bf/ppo_gridnet_diverse_encode_decode_sb3.py#L317-L321) how I did it:
* implement subclass of `Distribution` (this step is obvious)
* "replace" value of `self.action_dist` after instantiation of `ActorCriticPolicy` (subclass) -- this step feels ... weird, as `__init__` already made an initialization and reimplementation of entire `__init__` seems too intrusive from API perspective
* override `_get_action_dist_from_latent(Tensor) -> Distribution` ([here](https://github.com/kachayev/gym-microrts-paper-sb3/blob/48bcc7f609bb85eeffc5ae0dee5e64125abc59bf/ppo_gridnet_diverse_encode_decode_sb3.py#L336-L337)) as the method itself raises `ValueError` on any unknown distribution ([here](https://github.com/DLR-RM/stable-baselines3/blob/798b16aaf76ec051498a730a1d922e00184b4473/stable_baselines3/common/policies.py#L620)) -- this step feels weird for two reasons: 1) methods that starts from `_`, ideally, should be treated as implementation details for the class not public API intended for introduction of changes in the behavior; 2) raising `ValueError` seems unnecessary if probability distribution was already instantiated.
Note, that I also got lucky as the action space definition is known to `make_proba_distribution` helper. Otherwise replacing `self.action_dist` afterwords wouldn't work.
### Pitch
Honestly, I'm not sure what would be the best solution. It seems like with the current implementation of `ActorCriticPolicy` the entire notion of the "distribution of actions" is split between the policy itself and `mlp_extractor`: the policy decided on which distribution to use, `mlp_extractor` returns latent values that are further used by the policy to parametrized previously instantiated distribution. Which makes all of them intertwined in not very obvious way. It seems like we can simplify the flow by allowing `mlp_extractor` to return already parametrized distribution, so AC only needs to use already defined API of the distribution object (sample, log_prob, etc). In this case, `mlp_extractor` carries both sides coin: distribution + parametrization of the distribution.
Would love to hear thoughts and feedback on this. Also, would be happy to work on the PR when/if the direction is clear.
### Alternatives
TODO: fill this part in as the discussion goes.
### Checklist
- [x] I have checked that there is no similar [issue](https://github.com/DLR-RM/stable-baselines3/issues) in the repo (**required**)
Contributor guide
Assessment
This issue has not been assessed yet.