Add JAXsim Notebook for PD Controller with Gravity Compensation

[flferretti]

This PR will add a notebook with pixi support as a quickstart to the basic features of JAXsim.

In particular, the example will concern the creation of a PD controller with gravity compensation for a 2 DoFs cartpole system.

[flferretti]

For some reason, __cuda is not available for pixi when adding jaxlib with pixi add "jaxlib=*=*cuda*". Therefore I just dropped GPU support with pixi for the time being.

Moreover, I dropped GPU rendering for mujoco as it wasn't actually necessary for this example and it triggers an error when GPU is not available (as in pixi) setting MUJOCO_GL=egl

fyi @traversaro

[traversaro]

For some reason, __cuda is not available for pixi when adding jaxlib with pixi add "jaxlib=*=*cuda*". Therefore I just dropped GPU support with pixi for the time being.

Did you added the cuda = "11.8" under [system-requirements] group? See https://github.com/traversaro/mujoco-mjx-pixi-example/blob/7fc3d419f138c2114c0973b1d5810ac00dbff0dd/pixi.toml#L11 ? Indeed I thought this was documented in https://pixi.sh/configuration/#the-system-requirements-table, but perhaps in a bit cryptic way.

[traversaro]

Indeed I thought this was documented in https://pixi.sh/configuration/#the-system-requirements-table, but perhaps in a bit cryptic way.

Opened issue upstream on this: prefix-dev/pixi#590 .

[traversaro]

For some reason, __cuda is not available for pixi when adding jaxlib with pixi add "jaxlib=*=*cuda*". Therefore I just dropped GPU support with pixi for the time being.

Did you added the cuda = "11.8" under [system-requirements] group? See https://github.com/traversaro/mujoco-mjx-pixi-example/blob/7fc3d419f138c2114c0973b1d5810ac00dbff0dd/pixi.toml#L11 ? Indeed I thought this was documented in https://pixi.sh/configuration/#the-system-requirements-table, but perhaps in a bit cryptic way.

Anyhow, this is not blocking.

[traversaro]

Just to understand, how is the notebook supposed to be use now? pixi, Colab or both?

[flferretti]

Thanks Silvio, I'll give it a try. The example can currently run on both Colab and Jupyter using pixi

[flferretti]

The comment above was solved by changing the channels order in the pixi.toml as

< channels = ["nvidia", "conda-forge"]
> channels = ["conda-forge", "nvidia"]

Thus, this downloads cudatoolkit from conda-forge channel instead of from the nvidia channel, which does not provide ptxas:

XlaRuntimeError: FAILED_PRECONDITION: Couldn't get ptxas/nvlink version string: INTERNAL: Couldn't invoke ptxas --version

Which I solved by:

1. starting with a blank project with only the nvidia channel
2. pixi add cuda-nvcc
3. Add conda-forge in channels as: channels = ["conda-forge", "nvidia"]
4. pixi add "jaxlib=*=*cuda*"
5. Add the rest of the packages

C.C. @traversaro

[traversaro]

Thus, this downloads cudatoolkit from conda-forge channel instead of from the nvidia channel, which does not provide ptxas:

Cool! Hopefully once the CUDA 12 migration is completed, we can just install ptxas as well from conda-forge.

[traversaro]

Ok, just a small comment: at least of early 2024, you do not need to actually specify jaxlib=*=*cuda* to actually have jax with cuda enabled. As long as __cuda virtual package is available, the jaxlib with cuda support will be installed by default, see https://github.com/conda-forge/jaxlib-feedstock/blob/178711f83c5f7d13ccfb18c5989209f7afcfe78d/recipe/meta.yaml#L5 . Specifying jaxlib=*=*cuda* is probably safer as it will continue to install the cuda version even if in the future the default behaviour changes, but if we have problem to specify it at the moment in pixi, we can avoid it.

[traversaro]

Thanks @flferretti ! Just some minor comments (in the meanwhile we can start merging):

• In the notebook the cartpole.urdf still refers to my branch, it should refer instead to the master branch of jaxsim
• Is the notebook meant to be use with both pixi and colab? Could it make sense to have a mini README explaining how to launch it with pixi?
• Not sure how we can do, but we should make the install step (at least the one involving apt) dependent on whater the notebook runs on colab? When using pixi on a local machine, the apt commands will fail, I guess. The alternative (as done for ffmpeg) is to use bash to only apt install ffmpeg if the ffmpeg command is not available
• On the pixi side, perhaps we should install gazebo to ensure that rod's URDF support works fine?

P.s. do you have any idea how I can comment by line the notebook?

[flferretti]

Thanks @traversaro! For what regards:

• The wrong URL to the URDF, this was fixed in cff6337
• In ff74624 I added a README file to the examples folder
• The apt installation does not fail when using pixi, yet I followed your suggestion installing Gazebo only in Colab (see cc8f4f1)
• The current jaxlib=*=*cuda* works fine, so I would stick to that as it's currently the safest solution
• For reviewing notebooks we could use NB Review, which is free for GH Education. Yet, this uses a bot to post comments which may not be desirable, we could discuss this f2f

[traversaro]

Thanks @traversaro! For what regards:

* The wrong URL to the URDF, this was fixed in [cff6337](https://github.com/ami-iit/jaxsim/commit/cff6337ec20840d1d7661d98579628c4991be1df)

* In [ff74624](https://github.com/ami-iit/jaxsim/commit/ff74624bf593b89eeaa4bb8dc85bb9facde13b5b) I added a README file to the `examples` folder

* The apt installation does not fail when using `pixi`, yet I followed your suggestion installing Gazebo only in Colab (see [cc8f4f1](https://github.com/ami-iit/jaxsim/commit/cc8f4f193ccaac6913b2cff8f781264abfc65f8a))

* The current `jaxlib=*=*cuda*` works fine, so I would stick to that as it's currently the safest solution

* For reviewing notebooks we could use [NB Review](https://github.com/marketplace/review-notebook-app/plan/MDIyOk1hcmtldHBsYWNlTGlzdGluZ1BsYW4yMDIz#plan-2023), which is free for GH Education. Yet, this uses a bot to post comments which may not be desirable, we could discuss this f2f

Great, thanks! For me it is good to go.