Skip to content

Commit

Permalink
Updates installation instruction for SKRL JAX (#1570)
Browse files Browse the repository at this point in the history
# Description

SKRL provides examples using JAX, which by default only installs the CPU
version. To be compatible with CUDA, it is required to install JAX prior
to installing SKRL. This PR updates the installation instructions for
SKRL to note this requirement.

## Type of change

- This change requires a documentation update

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->
  • Loading branch information
kellyguo11 authored Dec 20, 2024
1 parent 052975a commit dd044a4
Showing 1 changed file with 6 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -122,11 +122,16 @@ SKRL
.. tab-item:: JAX

.. warning::

It is recommended to `install JAX <https://jax.readthedocs.io/en/latest/installation.html>`_ manually before proceeding to install skrl and its dependencies, as JAX installs its CPU version by default. For example, ``pip install -U "jax[cuda12]"`` can be used to install JAX for CUDA 12.
Visit the **skrl** `installation <https://skrl.readthedocs.io/en/latest/intro/installation.html>`_ page for more details.

.. code:: bash
# install python module (for skrl)
./isaaclab.sh -i skrl
# install skrl dependencies for JAX. Visit https://skrl.readthedocs.io/en/latest/intro/installation.html for more details
# install skrl dependencies for JAX
./isaaclab.sh -p -m pip install skrl["jax"]
# run script for training
./isaaclab.sh -p source/standalone/workflows/skrl/train.py --task Isaac-Reach-Franka-v0 --headless --ml_framework jax
Expand Down

0 comments on commit dd044a4

Please sign in to comment.