Rendering to the HTML doesn't work

[Denys88]

When I try to play saved html visualization on my local computer I always get this error in HTML:
Uncaught TypeError: Cannot read property '0' of undefined
at demo.html:191
at Array.map ()
at demo.html:191
at Array.forEach ()
at createAnimationClip (demo.html:189)
at demo.html:314

Everything works fine in the google colab.

[Denys88]

Some update:
I've found the root: you cannot visualize a batch of the agents.
And even with one agent I had shape: (1, 299). After fix it works.
But you definitely need to support multiple objects or at least show error.

[erikfrey]

Thanks for this feedback! Agreed, we should at least give a more informative error if batched trajectories are sent to the visualizer. We'll fix that and resolve this issue when it's posted.

Supporting visualizing batches of agents ( such as https://spectrum.ieee.org/image/MzMwNDQ2Ng.jpeg ) would be great too, but we would have to verify three.js can handle large scenes.

[Denys88]

Thanks, btw great job! I was able to integrate it to the pytorch and achieved pretty great training performance on 3090 :) more than 0.5m training steps per second and 14k scores with humanoid.

[erikfrey]

Very cool! Using pytorch, does each step have to copy back to CPU? Can pytorch access jax datastructures?

[Denys88]

`
def jax_to_torch(tensor):
from jax._src.dlpack import (to_dlpack,)
tensor = to_dlpack(tensor)
tensor = tpack.from_dlpack(tensor)
return tensor

def torch_to_jax(tensor):
from jax._src.dlpack import (from_dlpack,)
tensor = tpack.to_dlpack(tensor)
tensor = from_dlpack(tensor)
return tensor
`

I found a way. If use this code it works at least 10% faster compared to the CPU version via numpy.

[erikfrey]

Cool trick! If your project has a url, feel free to share it. We may eventually have a "Projects using Brax" section in the README.

Thanks for sharing!

[Denys88]

sounds good.
I just added it to my project: https://github.com/Denys88/rl_games/ but didn't update readme yet, will do it this evening.
Here is env link https://github.com/Denys88/rl_games/blob/master/rl_games/envs/brax.py
Actually I implemented training fully on GPU, only read vlaues to report it to the tensorfow.

[erikfrey]

OK, this is addressed in 170405e - thanks Denys.

[ViktorM]

This bug still is present. Rendering doesn't work with batched environments, even when it's a single env batch.

[btaba]

Thanks @ViktorM for pointing it out! We still don't support rendering batches, but the RuntimeError on invalid shapes was accidentally removed. Added it back