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():
@@ -23,4 +27,4 @@ def test_mpi_weighted_mean():
                 logger.logkv_mean(name, val)
         d2 = logger.dumpkvs()
         if comm.rank == 0:
-            assert d2 == correctval
+            assert d2 == correctval

setup.py

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