facebookresearch / mbrl-lib

Library for Model Based RL

MIT License

952 stars 154 forks source link

[Feature Request] Add option to use `functorch` for `BasicEnsemble` #162

Open luisenp opened 2 years ago

luisenp commented 2 years ago

🚀 Feature Request

Change BasicEnsemble to optionally use functorch.vmap.

Motivation and Pitch

Is your feature request related to a problem? Please describe.

BasicEnsemble lets the user provide arbitrary models, which are stacked together using a very naive loop-based implementation. We should be able to do this more efficiently now using functorch.

natolambert commented 2 years ago

Torch has a vmap??? !!!

luisenp commented 2 years ago

It's via functorch which is a separate library but I think they work closely with PyTorch's team.