Lightning-Universe/lightning-bolts

Integration with TorchRL

fedebotu opened this issue ยท 3 comments

๐Ÿš€ Feature

I think it would be great to integrate the current RL bolts models with the new TorchRL library

Motivation

While gym is a great library for environments, it may not be very flexible for certain applications and has limitations (such as the need to run environments on CPU). This is where the new PyTorch library torchrl can help with features such as:

  • Batched environments
  • TensorDicts are a great way to handle data
  • Can directly run the environments on single and multiple GPUs
  • Can wrap multiple environments, including gym
  • Convenient implementations of models and methods such as rollouts

Additional context

I uploaded in this repo an example notebook of the official TorchRL's pendulum tutorial with a custom environment that can run on batches and multiple GPUs

Borda commented

@fedebotu, that looks cool, would be interested in sending a PR with an integration proposal?