DLR-RM / DLR-RM/stable-baselines3

[Feature Request] API for using custom distributions with existing policies

Open
#703 10 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.