Dictionary Observations

[J-Travnik]

In machine learning, input comes in the form of matrices. Typically, models take in 1 matrix at a time, such as in the case of image classification where matrix containing the input image is given to the model and the model classifies the image.
However, there are many situations in which taking multiple inputs is necessary. One example is when training a reinforcement learning agent and the observations that the agent sees comes in the form of an image (e.g., camera, grid sensor, etc) and a vector describing the agent's state (e.g., current position, health, etc). In this situation, it is necessary to feed 2 inputs to the model. This PR addresses this.

Description

• added example environments with multi-input observations
• added DictReplayBuffer and DictRolloutBuffer to handle dictionaries
• added CombinedExtractor feature extractor that handles generic dictionary data
• added StackedObservations and StackedDictObservations to decouple data stacking from the VecFrameStack wrapper
• added test_dict_env.py test
• added a is_vectorized_env() method per observation space type in common\utils.py

Motivation and Context

• I have raised an issue to propose this change (link)

closes #216

closes #287 (image support for HER)
closes #284 (for off-policy algorithms)

Types of changes

• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)

Checklist:

• I've read the CONTRIBUTION guide (required)
• I have updated the changelog accordingly (required).
• My change requires a change to the documentation.
• I have updated the tests accordingly
• I have updated the documentation accordingly.
• I have reformatted the code using make format
• I have checked the codestyle using make check-codestyle and make lint
• I have ensured make pytest and make type both pass. (required)
• I have checked that the documentation builds using make doc (required)

TODOs

• check that documentation is properly updated
• check that dict with vectors only is the same as mlp policy + vector flatten
• Update env checker
• (optional) refactor HER: https://github.com/DLR-RM/stable-baselines3/tree/feat/refactor-her
• test A2C/PPO/SAC alone with GoalEnv

[Miffyli]

Quick comment before further commits: Make sure your formatter has line lengths set to 127, not 80. Seems like your formatter is splitting lines of allowed length (e.g. this one). black runs perfectly on windows with simple pip install black and then black -l 127.

[J-Travnik]

@araffin

After a quick look, I'm not sure we need that much complexity.
You could probably use the BitFlippingEnv together with a small dummy Custom Env.

Thats a good point though, I did want to make sure that it worked with Images though as that was my main use case. I think its important to cover this case (at least in a test) to make sure that VecEnvStacking and multiple environments all played well together when using images. Thoughts?

it would nice to have at least one common "combined extractor " for the Box.

I agree though I still see value in having the option to separate them. Perhaps this could come in the form of a bool with true default of "auto_concatenate_box_obs".

After a quick look at your PR, I think it would be better to almost duplicate the buffer classes, because having too many branches (with the if else) is hard to follow.

Hmm, that seems like a decent idea. Specifically you mean to create:

class DictReplayBuffer(ReplayBuffer):...

class DictRolloutBuffer(RolloutBuffer):...

?

I can run it on my side if you want and push the results. But it would be better for you to setup black and isort at least (you can take a look at the makefile for the details of the commands).

Yea, that'd be awesome. I will figure out black and stuff on my end. @Miffyli thanks for the pointer!

[araffin]

I think its important to cover this case (at least in a test) to make sure that VecEnvStacking and multiple environments all played well together when using images. Thoughts?

two things: if you want to test performances, then you need your environment (but we may skip the test if it is slow to run).
If you want only to check that the code run, you can use a dummy env that outputs dummy images like this one: https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/identity_env.py#L107
And of course, we want to check that it runs at least with different type of input (otherwise we could just use a "flatten" wrapper).

So, yes, we need to test performance to validate the implementation, but I would not run the test for each commit (maybe use pytest.mark.slow to skip it).

. Perhaps this could come in the form of a bool with true default of "auto_concatenate_box_obs".

sounds like a good compromise.

Hmm, that seems like a decent idea. Specifically you mean to create:

yes!

[Miffyli]

I agree though I still see value in having the option to separate them. Perhaps this could come in the form of a bool with true default of "auto_concatenate_box_obs".

sounds like a good compromise.

After bit of internal discussion we can leave this out completely for now. We can tell people to do stacking on the environment side (or alternatively provide some wrapper that concatenates arrays of specific keys together). Less messy code to worry about (e.g. images of different resolution).

Bit of a heads up: This update would be very desired (I would be using it actively, too!), but being a big change there will be bunch of reviewing and comments ^^.

[J-Travnik]

@Miffyli Yea, no worries. I'd rather take the time and to review and fix changes to figure out something that is stable.

And fair enough on the "auto_concatenate_box_obs". Potentially it can take the form of a VecEnv wrapper so that the tensors passed to the CombinedExtractor can be used readily.

[H-Park]

Tentative timeframe of this getting merged? I need this feature to start the literature grounded action masking PR

[Miffyli]

@H-Park We will first do first release (v1.0) properly before merging this, soooo most daring estimate is just before end of the year, but more likely in the beginning of January.

On preliminary glances this looks overall the way that things will work, so I would semi-reliably say you can base your work on how things work in @J-Travnik's branch :). Do note that I am not 100% sure action-space masking should be included in the main repo here or in contrib repo. A fresh discussion on topic with new papers/results would be appropriate.

[araffin]

@J-Travnik I reformated the code and corrected one typing issue in preprocessing, however, it complains about VecFrameStack now
I also fixed the check for HER, so the tests for HER should pass now ;)

@H-Park there is no deadline for this feature (at least after v1.0), it will take some time to have a working and clean solution.
But in the meantime, you could start working on @J-Travnik branch ;) (even though it may change).

EDIT: I did a first pass on the code, I will do another one once you have separated the buffers (quite hard to read for now)

[J-Travnik]

I added DictReplayBuffer and DictRolloutBuffer so that there is only one "if" statement in the policy to choose which one. They are child classes of their respective parents so the typing is preserved throughout the code base.
I'll go back and update the documentation for them probably early next week.

test_offpolicy_normalization fails now but only for HER and I am not completely sure why yet but I feel like there is an assumption of having ObsDictWrapper somewhere.

Going forward, what should the usage of the ObsDictWrapper be if dict observations can be handled by buffers?
I think there are some functions in the ObsDictWrapper that can be moved to a util script but in general I am not sure.

[araffin]

quick remark: black is run with a line length of 127 (cf contrib guide and PR template)
also, please use the latest version of it

EDIT: and we run isort too

[araffin]

@J-Travnik @Miffyli The refactoring of HER is done, tested and merged in here (I also added proper timeout handling).
We mostly need to do a final review of this PR and it is good to be merged =)

[JadenTravnik]

@araffin The integration of HER looks awesome. :)
I did a pass and didn't find much to comment on so koodos 👍

[Miffyli]

I think its worth doing the check here in the init. We have everything to make an informative error message so we might as well.

@J-Travnik sorry to bother one more time, but could you clarify what error message you meant here? is_image_space is checked inside the transpose_space function because it gets called from other places.

[Miffyli]

Looks positively brilliant to me! 👍 . If no further changes are required (see above comment), feel free to merge @araffin

@J-Travnik 🍻 / ☕ / 🍵 / hot-chocos on me if we ever get to meet in person! Thanks for all your time and patience! ^^

[araffin]

Finally merged 🎉
Thanks for the PR =D!

[J-Travnik]

@Miffyli @araffin This is wonderful :)
Looking forward to hot cocos in the future ;)