Turn on Github CI for coordinator_test (a rough approximation to running whole package E2E)

.github/workflows/core_test.yml

@@ -35,4 +35,4 @@ jobs:
           pip freeze
       - name: Test with pytest  # TODO(team): Fix tensorflow version conflict.
         run: |
-          # pytest -n auto iris
+          pytest -n auto iris/coordinator_test.py

iris/coordinator_test.py

@@ -19,12 +19,10 @@
 from typing import cast
 
 from absl.testing import absltest
-import gym
 from iris import checkpoint_util
 from iris import coordinator
 from iris.algorithms import ars_algorithm
-from iris.policies import nn_policy
-from iris.workers import rl_worker
+from iris.workers import simple_worker
 import launchpad as lp
 from ml_collections import config_dict
 import numpy as np
@@ -35,29 +33,6 @@
 _TEST_CHECKPOINT = "./testdata/test_checkpoint.pkl"
 
 
-class TestEnv(gym.Env):
-
-  def __init__(self):
-    self._ac_dim = 6
-    self._ob_dim = 14
-    self.action_space = gym.spaces.Box(
-        -1 * np.ones(self._ac_dim), np.ones(self._ac_dim), dtype=np.float32
-    )
-    self.observation_space = gym.spaces.Box(
-        -1 * np.ones(self._ob_dim), np.ones(self._ob_dim), dtype=np.float32
-    )
-
-  def step(self, action):
-    del action
-    return np.zeros(self._ob_dim), 1.0, False, {}
-
-  def reset(self):
-    return np.zeros(self._ob_dim)
-
-  def render(self, mode: str = "rgb_array"):
-    return np.zeros((16, 16))
-
-
 def make_bb_program(
     num_workers: int,
     num_eval_workers: int,
@@ -154,17 +129,14 @@ def setUp(self):
                     eval_rate=1,
                     num_iterations=400,
                     num_evals_per_suggestion=1,
-                    record_video_during_eval=True,
                 )
             ),
             worker=config_dict.ConfigDict(
                 dict(
-                    worker_class=rl_worker.RLWorker,
+                    worker_class=simple_worker.SimpleWorker,
                     worker_args=dict(
-                        env=TestEnv,
-                        policy=nn_policy.FullyConnectedNeuralNetworkPolicy,
-                        policy_args=dict(hidden_layer_sizes=[64, 64]),
-                        rollout_length=20,
+                        blackbox_function=np.sum,
+                        initial_params=np.zeros(shape=(10,)),
                     ),
                 )
             ),

requirements.txt

@@ -1,6 +1,9 @@
 typing  # Version dependent on Python version.
 absl-py>=1.0.0
-numpy>=1.21.5
+numpy>=1.21.5,<2.0.0
+
+jax>=0.4.34  # TODO:remove
+jaxlib>=0.4.34  # TODO:remove
 
 # Distributed systems libraries.
 # NOTE: Requires tensorflow~=2.8.0 to avoid proto issues.