[Feature] Weight Synchronization Schemes - Core Infrastructure

ghstack-source-id: 72840d1
Pull-Request: #3185

Author: vmoens

pyproject.toml

@@ -16,6 +16,7 @@ maintainers = [
 ]
 keywords = ["reinforcement-learning", "pytorch", "rl", "machine-learning"]
 classifiers = [
+    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
@@ -53,6 +54,7 @@ tests = [
     "pytest-benchmark",
     "pytest-rerunfailures",
     "pytest-error-for-skips",
+    "pytest-timeout",
 ]
 utils = [
     "tensorboard",
@@ -75,7 +77,7 @@ offline-data = [
 ]
 marl = ["vmas>=1.2.10", "pettingzoo>=1.24.1", "dm-meltingpot"]
 open_spiel = ["open_spiel>=1.5"]
-brax = ["jax[cuda12]>=0.7.0", "brax"]
+brax = ["jax>=0.7.0", "brax"]
 llm = [
     "transformers",
     "vllm",