TPU/GPU visualization for learning from pixels

[erikfrey]

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.

[masud99r]

Thanks for the great simulator. Any update on TPU/GPU visualization to enable learning from pixels?

[yardenas]

Hi!
I would be happy to try and help, any ideas where to start?
Could this or this be of any use?

I also like the idea of having a native jax solution as it would potentially enable differentiable rendering (analytical policy gradients from images, anyone?)

[erikfrey]

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!

[yardenas]

Thanks @erikfrey for the hints!

So just trying to get a better understanding: basically we would need to implement functions like create_capsule (e.g. here and its implementation here) which takes in the radius, half_height etc., but have it written in JAX, instead of C++, right?

[simon-bachhuber]

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

[JoeyTeng]

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.