Rendering to the HTML doesn't work about brax HOT 10 CLOSED

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.

from brax.

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.

from brax.

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.

from brax.

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

from brax.

`
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.

from brax.

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!

from brax.

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.

from brax.

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

from brax.

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

from brax.

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

from brax.