From 017e476f49b48938c5a9b5415417d0e0f6122c4e Mon Sep 17 00:00:00 2001
From: robcaulk <rob.caulk@gmail.com>
Date: Wed, 5 Oct 2022 17:20:40 +0200
Subject: [PATCH] add extras to setup.py for RL

---
 setup.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index b3693c9f9..304567bcc 100644
--- a/setup.py
+++ b/setup.py
@@ -15,6 +15,14 @@ freqai = [
     'scikit-learn',
     'catboost; platform_machine != "aarch64"',
     'lightgbm',
+    'xgboost'
+]
+
+freqai_rl = [
+    'torch',
+    'stable-baselines3',
+    'gym==0.21',
+    'sb3-contrib'
 ]
 
 develop = [
@@ -36,7 +44,7 @@ jupyter = [
     'nbconvert',
 ]
 
-all_extra = plot + develop + jupyter + hyperopt + freqai
+all_extra = plot + develop + jupyter + hyperopt + freqai + freqai_rl
 
 setup(
     tests_require=[
@@ -90,6 +98,7 @@ setup(
         'jupyter': jupyter,
         'hyperopt': hyperopt,
         'freqai': freqai,
+        'freqai_rl': freqai_rl,
         'all': all_extra,
     },
 )