mbrl-lib icon indicating copy to clipboard operation
mbrl-lib copied to clipboard
verified publisher icon facebookresearch

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

Open luisenp opened this issue 3 years ago • 2 comments

🚀 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.

luisenp avatar Aug 05 '22 15:08 luisenp

Torch has a vmap??? !!!

natolambert avatar Aug 05 '22 16:08 natolambert

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

luisenp avatar Aug 05 '22 16:08 luisenp