Rendering without notebooks

[namheegordonkim]

Hi all,

Is rendering without notebooks supported? I personally don't like using notebooks and am much more used to working with e.g. GLUT windows. It seems that wrapping brax environment into gym as per the third tutorial (https://colab.research.google.com/github/google/brax/blob/main/notebooks/pytorch.ipynb) doesn't let me call env.render() (throws NotImplementedError).

Any thoughts would be highly appreciated. Thanks!

[kayuksel]

Yes, that would be a great feature, even if that means not rendering whole batch but a selected one. If we could render all of them in a single scene (where they do not collide), that could also be good.

[bayerj]

I have been using streamlit for showing the renderings. If you have the html, it is as easy as this:

import streamlit.components.v1 as components
html_string = load_html_rendering()    # this would have to be dumped in a different process
components.html(html_string, height=500)

Then reloading the streamlit app shows it. That's also nice to show all kinds of other analytics. Here is a screenshot of an app I built:

[erikfrey]

We have begun work on this. No ETA yet but it's an active priority for us.

[EelcoHoogendoorn]

Interesting; that streamlit thing looks pretty slick.

I was running into the same issues myself; not very familiar with 3JS and I think itd be cleaner to keep the visualization within python. Or at least have multiple rendering options available.

The best thing I can come up with would be to use the high level vispy/scene module; which is a similar high level API to 3JS; though its a bit stale in terms of maintenance... but I suppose its fine for putting some cylinders and boxes on screen. Not sure if there is a better option in the python ecosystem; doing it in raw pyopengl sure would be a kludge.

[erikfrey]

OK, rendering to pixels is now available, but not yet documented. If you'd like to try it, please check out:

https://github.com/google/brax/blob/main/brax/io/image.py

Please keep in mind this is CPU rendering, so better for eval rendering and other programmatic use cases, rather than training. We will move to GPU/TPU rendering in the future, which should be suitable for training.

In the coming days we'll update our colabs with an example of how to use it, and also make env.render return rgb arrays using this module.

[erikfrey]

An example of how to use it programmatic rendering is now live at notebooks/environments.ipynb

[EelcoHoogendoorn]

Thanks @erikfrey ! Ive also found use in the programmatic rendering, ive found it a better debug workflow for me in making changes to the physics than having to switch to notebooks and restarting those.

Ive been implementing a vispy based rendering as well; partially as an exercise, but also because im more familiar with it. It a pretty clean 200 lines of code and as far as I can tell its feature-complete in the sense that it works for all envs i tested it on. One thing I like about vispy is that it allows you to drop into raw opengl if required; I can see that being nice if wanting to visualize the rollout of a large batch efficiently using instancing if that becomes useful. Also it works both to generate static images as well as windowed interaction (or in notebooks too i think tho i havnt tested that), and it uses the GPU if available, obviously.

However, not sure what to do with this code. Since I cant/wont promise maintenance with future updates myself perhaps I could create a PR for this inside the experimental submodule?

[erikfrey]

@EelcoHoogendoorn that's awesome! Offhand do you know how fast it renders on CPU only? Can you open a new issue and point us to your example code using vispy? If it works well, I think this would be a candidate for Brax to support, as GPU rendering is on our roadmap.

[EelcoHoogendoorn]

Hmm, just looking at the contributors license, I am reminded I should check this with my employer. Not that this has anything to do with my day job, but they havn't been particularly clear about what they think is ok and what isn't. Hopefully ill get some clarity on that soon.

Didnt test on CPU yet. Given the simplicity of the scenes, I expect it to be fine. In my experience in combining opengl and python, the hard part is not getting bogged down by python overhead. Im currently precomputing all qp->4x4 logic, so an animation loop consist of just updating those 4x4s in the scenegraph, and one draw call per body in the system. I wont vouch for how performant even that will be in python if you are aiming for 1000s of independent non-instanced bodies, but I dont expect fill rate or anything like that to be your biggest enemy, on a modern cpu, with scenes of more typical complexity.

[EelcoHoogendoorn]

Didnt spend any time on styling or prettyness yet, but fwiw, here is a torus getting a shove...

[EelcoHoogendoorn]

Speaking of which, I see that my manager is also on github. Hi @amoshaviv, what are our OS contribution policies like nowadays?

[yardenas]

@erikfrey what's the timeline on the GPU rendering project? I might be able to contribute some parts.