Fix python compatibility

setup.py

@@ -1,8 +1,8 @@
 from setuptools import setup
 import sys
 
-assert sys.version_info.major == 3 and sys.version_info.minor >= 6, \
-    "This repo is designed to work with Python 3.6 and greater." \
+assert sys.version_info.major == 3 and sys.version_info.minor in [6, 7] , \
+    "This repo is designed to work with Python 3.6 or 3.7." \
     + "Please install it before proceeding."
 
 setup(
@@ -23,7 +23,6 @@
         'sklearn',
         'imageio',
         'seaborn==0.8.1',
-        'tensorflow-gpu<2',
         'dm-sonnet<2',
         'tensorflow-probability==0.7.0',
         'torch==1.4.0',
@@ -39,3 +38,18 @@
     description="TeachMyAgent: A benchmark to study and compare ACL algorithms for DeepRL in continuous procedural environments.",
     author="Clément Romac",
 )
+
+# ensure there is some tensorflow build with version above 1.4
+import pkg_resources
+import re
+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 1.4'
+from distutils.version import LooseVersion
+tf_version = LooseVersion(re.sub(r'-?rc\d+$', '', tf_pkg.version))
+assert tf_version >= LooseVersion('1.4.0') and tf_version <= LooseVersion('1.15.5'), \
+    'TensorFlow version between 1.4 and 1.15.5 required'