can `Brax` simulate on the GPU (what args to use?)

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.

Apache License 2.0

2.25k stars 249 forks source link

Closed Kallinteris-Andreas closed 1 year ago

Kallinteris-Andreas commented 1 year ago

Looking at the reference environments (e.g.) https://github.com/google/brax/blob/main/brax/envs/half_cheetah.py there is no device argument, it is possible to run the environment on a GPU or TPU?

Or did I read the README.md incorrectly

erikfrey commented 1 year ago

Hello, you can specify a default device in JAX or specify it at time of jit.