Unverified Commit dd044a49 authored by Kelly Guo's avatar Kelly Guo Committed by GitHub

Updates installation instruction for SKRL JAX (#1570)

# 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
-->
parent 052975a9
......@@ -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
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment