Support Google Cloud TPUs

[andrewsykim]

I'd like to try using SkyRL with Google Cloud TPUs on GKE. Is it supported or possible to run TPUs today?

I also saw a recent issue to integrate Ray Direct Transport into SkyRL. We plan to add TPU support in RDT in the near future as well.

I'm happy to help contribute in this area with some guidance from maintainers :)

[pcmoritz]

It would also be wonderful to optimize the jax backend in https://github.com/NovaSky-AI/SkyRL/tree/main/skyrl-tx better for TPUs. I ran some earlier versions of the code on TPUs but currently don't have TPU access, I'd love to help with benchmarking / optimizing if given access.