Does the visualizer need to depend on JAX?

[stephane-caron]

I went through installing dependencies manually, but my overall question for this project would be: couldn't the visualizer run with only MeshCat as dependency?

[stephane-caron]

(Verbose user feedback on!) After installing dependencies manually, I ran into the following import error:

RuntimeError: jaxlib is version 0.3.2, but this version of jax requires version >= 0.3.14.

I upgraded jaxlib, but then run into the following:

RuntimeError: jaxlib version 0.3.20 is newer than and incompatible with jax version 0.3.15. Please update your jax and/or jaxlib packages.

So I guess it will be jaxlib==0.3.15 for now. But given other constraints it won't be able to stay at that precise version for long.

[diegoferigo]

I went through installing dependencies manually, but my overall question for this project would be: couldn't the visualizer run with only MeshCat as dependency?

Short answer: yes. Long answer here below.

The visualizer currently needs two piece of information from jaxsim:

1. The structure of the kinematic tree. This is necessary to build the node tree used by meshcat. I used jaxsim.parser.kinematic_graph.KinematicGraph since it was already done, but I plan to move this functionality to ami-iit/rod as soon as rod becomes the defaul parser of jaxsim.
2. Computing forward kinematics. This is necessary to compute the transforms of all links, and again I used jaxsim because the algorithm was already implemented there. I agree with you that it would be nice porting the FK algo here in meshcat-viz so that the jaxsim dependency could be dropped for good. Also this is planned, even if with lower priority.

For what regards the dependency of JAX, I'm still debating. I'm not sure what could be the performance achievable with pure numpy, and having a JIT-compiled FK algorithm in JAX (even just CPU) could be quite interesting.

Regarding the versioning problems of jax/jaxlib, it's a consequence of what described in ami-iit/jaxsim#12. Sadly JAX deprecated and removed some experimental module that jaxsim was using, and I still need to finalize the activity to switch to other resources as documented in ami-iit/jaxsim#12.

[diegoferigo]

Btw, I reproduced the problem of jax/jaxlib versions in https://github.com/ami-iit/jaxsim/actions/runs/3183397287/jobs/5190569189

[diegoferigo]

I made some progress but got stuck again (ami-iit/jaxsim#23). I guess it's time to migrate the jaxsim parser to rod.

[diegoferigo]

Btw, I reproduced the problem of jax/jaxlib versions in https://github.com/ami-iit/jaxsim/actions/runs/3183397287/jobs/5190569189

I isolated the continuous integration part of ami-iit/jaxsim#23 in ami-iit/jaxsim#24. This PR also fixes the jaxlib version problem described in #4 (comment).

[diegoferigo]

In #7, I removed the direct dependency from jaxsim. We use external libraries to compute FK, and now it can be done either with jaxsim or idyntree (check meshcat_viz.fk).