From 713f1a0aec6981dbab7cd2243a2ba1e69b337123 Mon Sep 17 00:00:00 2001
From: Isaac Lascasas <isaaclascasas@gmail.com>
Date: Sat, 26 Oct 2019 00:50:04 +0200
Subject: [PATCH] tf2: Updated setup.py dependencies. (#1002)

* Updated setup.py dependencies.

* Allow tf2 baselines pip package setup to work with all tf2 cpu/gpu versions.
---
 setup.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/setup.py b/setup.py
index 4d19d7900d..b549319de3 100644
--- a/setup.py
+++ b/setup.py
@@ -12,10 +12,10 @@
         'filelock',
         'pytest',
         'pytest-forked',
-        'atari-py'
+        'atari-py~=0.2.0'
     ],
     'bullet': [
-        'pybullet',
+        'pybullet'
     ],
     'mpi': [
         'mpi4py'
@@ -32,8 +32,7 @@
       packages=[package for package in find_packages()
                 if package.startswith('baselines')],
       install_requires=[
-          'tf-nightly-2.0-preview',
-          'gym',
+          'gym<0.14',
           'scipy',
           'tqdm',
           'joblib',
@@ -50,3 +49,14 @@
       author_email='gym@openai.com',
       version='0.1.5')
 
+# ensure there is some tensorflow build with version above 2.0
+import pkg_resources
+tf_pkg = None
+for tf_pkg_name in ['tensorflow', 'tensorflow-gpu', 'tf-nightly', 'tf-nightly-gpu']:
+    try:
+        tf_pkg = pkg_resources.get_distribution(tf_pkg_name)
+    except pkg_resources.DistributionNotFound:
+        pass
+assert tf_pkg is not None, 'TensorFlow needed, of version above 2.0'
+from distutils.version import LooseVersion
+assert LooseVersion(re.sub(r'-?rc\d+$', '', tf_pkg.version)) >= LooseVersion('2.0.0')
\ No newline at end of file