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

@@ -12,10 +12,10 @@ extras = {
         'filelock',
         'pytest',
         'pytest-forked',
-        'atari-py'
+        'atari-py~=0.2.0'
     ],
     'bullet': [
-        'pybullet',
+        'pybullet'
     ],
     'mpi': [
         'mpi4py'
@@ -32,8 +32,7 @@ setup(name='baselines',
       packages=[package for package in find_packages()
                 if package.startswith('baselines')],
       install_requires=[
-          'tf-nightly-2.0-preview',
+          'gym<0.14',
-          'gym',
           'scipy',
           'tqdm',
           'joblib',
@@ -50,3 +49,14 @@ setup(name='baselines',
       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')