fix: include HTML and JS assets in wheel

[cemlyn007]

At the momentvisualizer is not in the wheel, this causes an error when you want to render html using brax.v2.io.html.render_from_json which by default makes use of HTML and JS assets found within the visualizer folder.

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

env = brax.v2.envs.ant.Ant()
brax.v2.io.html.render(env.sys, [env.reset(jax.random.PRNGKey(0)).pipeline_state])

Raises:

FileNotFoundError: [Errno 2] No such file or directory: '<REDACTED>/venv/lib/python3.10/site-packages/brax/v2/visualizer/index.html'

Because of:

html_path = epath.resource_path('brax') / 'v2/visualizer/index.html'

Which is not contained in the wheel.

Now a simple solution would be to include the files in visualizer in the wheel but I think you guys should consider moving the HTML and JS assets to possibly io because currently io and visualizer depend on each other, whether for Python code or HTML and JS assets.

I am happy to make an MR that moves the assets to brax.v2.io and make sure the assets are added to the distributable but please let me know what you would like to do.

Have a good weekend,

Cemlyn

[btaba]

Hi @cemlyn007 thanks for pointing this out! We added those files to the manifest so should be ok now. Let us know if there are other issues