feat: uv (#1171)

* wip: switch to uv completely

* chore: move python version to uv setup

* chore: revert to uv pip for CI

* chore: revert integration test rule used for testing

* chore: add max python version and setuptool to requirements

* chore: add line to activate venv in readme

* chore: update detailed install

* chore: pre-commit

Author: sash-a

.github/workflows/ci.yaml

@@ -17,18 +17,19 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v3
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.4"
+          version: "0.6.12"
           enable-cache: true
-          cache-dependency-glob: "requirements/requirements**.txt"  # invalidate cache when requirements file changes
+          cache-dependency-glob: "uv.lock"  # invalidate cache when requirements file changes
+          python-version: "${{ matrix.python-version }}"
 
       - uses: actions/setup-python@v5
         with:
             python-version: "${{ matrix.python-version }}"
 
       - name: Install python dependencies 🔧
-        run: uv pip install .[dev]
+        run: uv pip install . --group dev
         env:
           UV_SYSTEM_PYTHON: 1
 

.github/workflows/integration_tests.yaml

@@ -20,18 +20,19 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v3
+        uses: astral-sh/setup-uv@v5
         with:
-          version: "0.5.4"
+          version: "0.6.12"
           enable-cache: true
-          cache-dependency-glob: "requirements/requirements**.txt"  # invalidate cache when requirements file changes
+          cache-dependency-glob: "uv.lock"  # invalidate cache when requirements file changes
+          python-version: "${{ matrix.python-version }}"
 
       - uses: actions/setup-python@v5
         with:
             python-version: "${{ matrix.python-version }}"
 
       - name: Install python dependencies 🔧
-        run: uv pip install .[dev]
+        run: uv pip install . --group dev
         env:
           UV_SYSTEM_PYTHON: 1
 

.gitignore

@@ -174,4 +174,6 @@ results/
 outputs/
 multirun/
 .neptune
-.aider*
+
+# AI
+.aider**

README.md

@@ -41,11 +41,20 @@ Mava allows researchers to experiment with multi-agent reinforcement learning (M
 
 ## Installation 🎬
 
-At the moment Mava is not meant to be installed as a library, but rather to be used as a research tool. We recommend cloning the Mava repo and pip installing as follows:
+At the moment Mava is not meant to be installed as a library, but rather to be used as a research tool. We recommend cloning the Mava repo and installing dependencies using [uv](https://github.com/astral-sh/uv) as follows:
 
 ```bash
-git clone https://github.com/instadeepai/mava.git
-cd mava
+# Clone the repository
+git clone https://github.com/instadeepai/Mava.git
+cd Mava
+# Create a virtual environment and install all dependencies
+uv sync
+# Activate the virtual environment
+source .venv/bin/activate
+```
+
+Alternatively with pip, create a virtual environment and then:
+```bash
 pip install -e .
 ```
 

docs/DETAILED_INSTALL.md

@@ -16,9 +16,7 @@ cd mava
 
 3. Create and activate a virtual environment and install requirements
 ```bash
-uv venv -p=3.12
-source .venv/bin/activate
-uv pip install -e .
+uv sync -p=3.12
 ```
 
 4. Install jax on your accelerator. The example below is for an NVIDIA GPU, please the [official install guide](https://github.com/google/jax#installation) for other accelerators.
@@ -29,6 +27,11 @@ uv pip install "jax[cuda12]==0.4.30"
 
 5. Run a system!
 ```bash
+uv run mava/systems/ppo/anakin/ff_ippo.py env=rware
+```
+or
+```
+source .venv/bin/activate
 python mava/systems/ppo/anakin/ff_ippo.py env=rware
 ```
 

pyproject.toml

@@ -8,11 +8,11 @@ include=['mava*']
 [project]
 name="id-mava"
 authors=[{name="InstaDeep Ltd"}]
-dynamic=["version", "dependencies", "optional-dependencies"]
+dynamic=["version"]
 license={file="LICENSE"}
 description="Distributed Multi-Agent Reinforcement Learning in JAX."
 readme ="README.md"
-requires-python=">=3.10"
+requires-python=">=3.10,<3.13"
 keywords=["multi-agent", "reinforcement learning", "python", "jax", "anakin", "sebulba"]
 classifiers=[
     "Environment :: Console",
@@ -28,10 +28,47 @@ classifiers=[
     "Topic :: Software Development :: Libraries :: Python Modules",
     "License :: OSI Approved :: Apache Software License",
 ]
+dependencies = [
+    "brax==0.10.3",
+    "colorama",
+    "distrax",
+    "flashbax~=0.1.0",
+    "flax>=0.8.1,<0.10.4",
+    "gigastep @ git+https://github.com/mlech26l/gigastep",
+    "gymnasium==1.0.0",
+    "hydra-core==1.3.2",
+    "id-marl-eval @ git+https://github.com/instadeepai/marl-eval",
+    "jax==0.4.30",
+    "jaxlib==0.4.30",
+    "jaxmarl @ git+https://github.com/RuanJohn/JaxMARL@unpin-jax", # This only unpins the version of Jax.
+    "jumanji>= 1.1.0",
+    "lbforaging",
+    "matrax>= 0.0.5",
+    "mujoco==3.1.3",
+    "mujoco-mjx==3.1.3",
+    "neptune",
+    "numpy==1.26.4",
+    "omegaconf",
+    "optax",
+    "protobuf~=3.20",
+    "rware",
+    "scipy==1.12.0",
+    "setuptools",
+    "smaclite @ git+https://github.com/uoe-agents/smaclite.git",
+    "tensorboard_logger",
+    "tensorflow_probability",
+    "type_enforced", # needed because gigastep is missing this dependency
+]
+
+[dependency-groups]
+dev = [
+    "mypy",
+    "pre-commit",
+    "pytest",
+]
+
 [tool.setuptools.dynamic]
 version={attr="mava.__version__"}
-dependencies={file="requirements/requirements.txt"}
-optional-dependencies={dev={file=["requirements/requirements-dev.txt"]}}
 
 [project.urls]
 "Homepage"="https://github.com/instadeep/Mava"