Re-add tf proba

araffin · Jun 15, 2023 · bbec22d · bbec22d
1 parent f0ee601
commit bbec22d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -38,6 +38,7 @@ openrlbenchmark = {version = "^0.1.1b4", optional = true}
 jax = {version = "^0.3.17", optional = true}
 jaxlib = {version = "^0.3.15", optional = true}
 flax = {version = "^0.6.0", optional = true}
+tensorflow-probability = {version = "^0.19.0", optional = true}
 optuna = {version = "^3.0.1", optional = true}
 optuna-dashboard = {version = "^0.7.2", optional = true}
 envpool = {version = "^0.6.4", optional = true}
@@ -70,7 +71,7 @@ plot = ["pandas", "seaborn"]
 pytest = ["pytest"]
 mujoco = ["mujoco", "imageio"]
 mujoco_py = ["free-mujoco-py"]
-jax = ["jax", "jaxlib", "flax"]
+jax = ["jax", "jaxlib", "flax", "tensorflow-probability"]
 docs = ["mkdocs-material", "markdown-include", "openrlbenchmark"]
 envpool = ["envpool"]
 optuna = ["optuna", "optuna-dashboard"]