fix tests - add matplotlib to setup_requires, put mpi4py import in try-except

baselines/common/test_mpi_util.py

@@ -1,7 +1,11 @@
 from baselines.common import mpi_util
-from mpi4py import MPI
 from baselines import logger
 from baselines.common.tests.test_with_mpi import with_mpi
+try:
+    from mpi4py import MPI
+except ImportError:
+    MPI = None
+    
 
 @with_mpi()
 def test_mpi_weighted_mean():

setup.py

@@ -12,7 +12,8 @@ extras = {
         'filelock',
         'pytest',
         'pytest-forked',
-        'atari-py'
+        'atari-py',
+        'matplotlib',
     ],
     'mpi': [
         'mpi4py'