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 170405e4880ea780d38ee5bba06930a74cb99832 - 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