ci/runtests.sh - pass all folders to pytest (#342)

* ci/runtests.sh - pass all folders to pytest

* mpi_optimizer_test precision 1e-4

* fixes to tests

* search for tests in the entire jax folder, also remove unnecessary humor

baselines/common/tests/test_with_mpi.py

@@ -4,6 +4,7 @@ import subprocess
 import cloudpickle
 import base64
 import pytest
+from functools import wraps
 
 try:
     from mpi4py import MPI
@@ -12,6 +13,7 @@ except ImportError:
 
 def with_mpi(nproc=2, timeout=30, skip_if_no_mpi=True):
     def outer_thunk(fn):
+        @wraps(fn)
         def thunk(*args, **kwargs):
             serialized_fn = base64.b64encode(cloudpickle.dumps(lambda: fn(*args, **kwargs)))
             subprocess.check_call([