search

alex-petrenko / sample-factory

High throughput synchronous and asynchronous reinforcement learning
https://samplefactory.dev
MIT License
811 stars 109 forks source link

TPU Support #277

Open subho406 opened 1 year ago

subho406 commented 1 year ago

This is an amazing library! I was wondering if the library has any support for TPUs? If not how difficult will it get it working in TPUs? I would be willing to contribute to this if its not super time consuming.

alex-petrenko commented 1 year ago

Not sure actually, never had a chance to work with TPUs.

I imagine it would be trivial for serial training usecase --serial_mode=True, just replacing "cuda" with TPU everywhere.

Although SF2 relies on shared memory tensors a lot. If TPU shared tensors behave the same way as CUDA tensors, then I assume it would be easy to implement.