Open erikfrey opened 2 years ago
Thanks for the great simulator. Any update on TPU/GPU visualization to enable learning from pixels?
Hi @yardenas and @masud99r - JAX rendering is still an area of active interest for us! One of those libraries might be a good starting point, although you'll also need to generate meshes from our brax primitives like capsules.
You can find examples of how to do that in our current cpu-only renderer code: https://github.com/erwincoumans/tinyrenderer
If you get on-device rendering going via JAX, please share with us a colab so we can check it out!
Brax's rendering is CPU-only at the moment. For agents with vision, it would be useful for rendering to happen on accelerator so that training can stay fast.
Vispy may be useful, but a fast native jax renderer would be even better for data on device compatability.
I am currently looking in this direction for my personal projects. I made a post there how one might do this efficiently. Might be interesting here too. Issue over at VisPy
Hi, I have created a pure JAX implementation and an adapter layer mimicking the behaviour of existing pytinyrenderer. The code will be released under Apache-2.0 License here. I will be working on a PR to replace current visualisation code in Brax soon.
A working example can be found under the root directory, e.g., this. A Colab is also available now, which adopts exact examples of the existing pytinyrenderer's.
Update: I have opened a draft PR #367 . There is still minor issue (plane may not be rendered correctly; performance issue), but it is working now. I will address the plane rendering issue first, then performance. The rendering issue is fixed. I'll be working on the performance improvements.
Brax's rendering is CPU-only at the moment. For agents with vision, it would be useful for rendering to happen on accelerator so that training can stay fast.
Vispy may be useful, but a fast native jax renderer would be even better for data on device compatability.