Issue with HTML saving since latest commit

[Theo-Cheynel]

I'm having issues with brax.io.html.save with the latest commits and the switch to v2.

if I checkout on 1d5fbd7cdc71708e8549e34a0aa485e7dc8a460c, and run the following code, I get a near-empty HTML, which displays as a blank page in any browser.

from brax import envs
from brax.io import html
import jax

env = envs.get_environment(env_name='ant', backend='positional')
state = jax.jit(env.reset)(rng=jax.random.PRNGKey(seed=0))
html.save('test.html', env.sys, [state.pipeline_state])

Expected behaviour : on the previous commit (fc8f410583e56b6ff38ddad0650bccfb80937e46), if I run the following code (same code, imports were just switched to v2), I get an HTML file with the 3d scene and the ant.

from brax.v2 import envs
from brax.v2.io import html
import jax

env = envs.get_environment(env_name='ant', backend='positional')
state = jax.jit(env.reset)(rng=jax.random.PRNGKey(seed=0))
html.save('test.html', env.sys, [state.pipeline_state])

Am I missing something about a new way to save HTML visualization files ?

[btaba]

This will be fixed once we do a release (very shortly), for now you can manually edit the line here: https://github.com/google/brax/blob/1d5fbd7cdc71708e8549e34a0aa485e7dc8a460c/brax/io/html.py#L45

to point to the recent commit https://cdn.jsdelivr.net/gh/google/brax@1d5fbd7/brax/visualizer/js/viewer.js

[Theo-Cheynel]

It works great ! Thanks :smile: (wanted to modify the HTML template to show the associated observation vector and rewards under the ThreeJS frame)