Repository: rainx0r/metaworld-algorithms
Issue: Add specific jax version to pyproject.toml file
Root cause: Configuration Error

When starting with a fresh clone and environment sync, i.e.

1. Create a virtual environment for the project: `uv venv .venv --python 3.12`
2. Activate the virtual environment: `source .venv/bin/activate`
3. Install the dependencies: `uv pip install -e .`

An error happens:
```
Traceback (most recent call last):
  File "/Users/atodrv/Documents/metaworld-algorithms/examples/multi_task/ppo_mt10.py", line 15, in <module>
    from metaworld_algorithms.rl.algorithms import PPOConfig
  File "/Users/atodrv/Documents/metaworld-algorithms/metaworld_algorithms/rl/algorithms/__init__.py", line 3, in <module>
    from .base import Algorithm, OffPolicyAlgorithm
  File "/Users/atodrv/Documents/metaworld-algorithms/metaworld_algorithms/rl/algorithms/base.py", line 8, in <module>
    import orbax.checkpoint as ocp
  File "/Users/atodrv/Documents/metaworld-algorithms/.venv/lib/python3.12/site-packages/orbax/checkpoint/__init__.py", line 23, in <module>
    from orbax.checkpoint import aggregate_handlers
  File "/Users/atodrv/Documents/metaworld-algorithms/.venv/lib/python3.12/site-packages/orbax/checkpoint/aggregate_handlers.py", line 24, in <module>
    from orbax.checkpoint._src.metadata import tree as tree_metadata
  File "/Users/atodrv/Documents/metaworld-algorithms/.venv/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/tree.py", line 36, in <module>
    from orbax.checkpoint._src.metadata import tree_rich_types
  File "/Users/atodrv/Documents/metaworld-algorithms/.venv/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/tree_rich_types.py", line 25, in <module>
    from orbax.checkpoint._src.metadata import value_metadata_entry
  File "/Users/atodrv/Documents/metaworld-algorithms/.venv/lib/python3.12/site-packages/orbax/checkpoint/_src/metadata/value_metadata_entry.py", line 25, in <module>
    from orbax.checkpoint._src.serialization import types
  File "/Users/atodrv/Documents/metaworld-algorithms/.venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/types.py", line 33, in <module>
    from orbax.checkpoint._src.serialization import serialization
  File "/Users/atodrv/Documents/metaworld-algorithms/.venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/serialization.py", line 403, in <module>
    dll: Optional[layout.DeviceLocalLayout],
                  ^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'jax.experimental.layout' has no attribute 'DeviceLocalLayout'
```


This is because `jax==0.7.0` and `jaxlib==0.7.0` are installed by default (`>=0.5.0`), while `orbax` requires `jax==0.6.0`. This is either fixable through:
1. Do `uv pip install jax==0.6.0` (or alternatively `uv pip install jax[cuda12]==0.6.0`
2. Add specific `jax` versions to `pyproject.toml`:
```
[project.optional-dependencies]
cpu = ["jax==0.6.0"]
# https://github.com/jax-ml/jax/issues/27062
metal = ["jax==0.6.0", "jax-metal>=0.1.0; sys_platform == 'darwin'"]
cuda12 = ["jax[cuda12]==0.6.0"]
tpu = ["jax[tpu]==0.6.0"]
testing = [
    "chex>=0.1.89",
    "pytest>=8.3.5",
]
```

This would work for `uv pip install -e ".[cpu]"`, `uv pip install -e ".[cuda12]"`, etc., but not for `uv pip install -e .`, unless the main dependencies are updated with `jax==0.6.0`:
```
dependencies = [
    "gymnasium>=1.0.0",
    "metaworld",
    "distrax>=0.1.5",
    "flax>=0.9.0",
    "jax==0.6.0",
    "jaxtyping>=0.2.34",
    "numpy>=2.1.2",
    "orbax-checkpoint<=0.11.11",
    "tyro>=0.8.11",
    "wandb<=0.19.9",
    "scipy>=1.15.2",
]
```




Hi, thank you for the suggestion.

I'm not entirely sure I want / expect `uv pip install -e .` to work. Putting jax into the main dependency list with no extras makes it install the CPU-only version as specified [here](https://docs.jax.dev/en/latest/installation.html#installation). I think a lot of people might end up just installing the package that way then complain about slow performance since it's not using their accelerator. So similar to JAX where you explicitly have to specify which accelerator you want to use, I think it might be a good idea to force users here to explicitly think about the accelerator they want to use and to install the package with the right extras enabled.

I think the README instructions are quite clear enough about how you're meant to use the package, but I'll add a small section regarding this for further clarity.

Ah, I see now that I have JAX>=0.5 on the extra dependencies which could be an issue. I'll pin it to 0.6.0 to ensure orbax compatibility (although `uv.lock` has specific JAX versions that should be honoured by `uv pip install`).