[RLlib] Tf2.x native. (#8752)

rllib/utils/exploration/epsilon_greedy.py

@@ -111,9 +111,13 @@ class EpsilonGreedy(Exploration):
             ),
             false_fn=lambda: exploit_action)
 
-        assign_op = tf1.assign(self.last_timestep, timestep)
-        with tf1.control_dependencies([assign_op]):
+        if self.framework in ["tf2", "tfe"]:
+            self.last_timestep = timestep
             return action, tf.zeros_like(action, dtype=tf.float32)
+        else:
+            assign_op = tf1.assign(self.last_timestep, timestep)
+            with tf1.control_dependencies([assign_op]):
+                return action, tf.zeros_like(action, dtype=tf.float32)
 
     def _get_torch_exploration_action(self, q_values, explore, timestep):
         """Torch method to produce an epsilon exploration action.

rllib/utils/exploration/gaussian_noise.py

@@ -72,7 +72,7 @@ class GaussianNoise(Exploration):
             0, framework=self.framework, tf_name="timestep")
 
         # Build the tf-info-op.
-        if self.framework == "tf":
+        if self.framework in ["tf", "tfe"]:
             self._tf_info_op = self.get_info()
 
     @override(Exploration)
@@ -123,11 +123,18 @@ class GaussianNoise(Exploration):
         logp = tf.zeros(shape=(batch_size,), dtype=tf.float32)
 
         # Increment `last_timestep` by 1 (or set to `timestep`).
-        assign_op = (
-            tf1.assign_add(self.last_timestep, 1) if timestep is None else
-            tf1.assign(self.last_timestep, timestep))
-        with tf1.control_dependencies([assign_op]):
+        if self.framework in ["tf2", "tfe"]:
+            if timestep is None:
+                self.last_timestep.assign_add(1)
+            else:
+                self.last_timestep.assign(timestep)
             return action, logp
+        else:
+            assign_op = (
+                tf1.assign_add(self.last_timestep, 1) if timestep is None else
+                tf1.assign(self.last_timestep, timestep))
+            with tf1.control_dependencies([assign_op]):
+                return action, logp
 
     def _get_torch_exploration_action(self, action_dist, explore, timestep):
         # Set last timestep or (if not given) increase by one.

rllib/utils/exploration/ornstein_uhlenbeck_noise.py

@@ -95,7 +95,11 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
             shape=[self.action_space.low.size], stddev=self.stddev)
         ou_new = self.ou_theta * -self.ou_state + \
             self.ou_sigma * gaussian_sample
-        ou_state_new = tf1.assign_add(self.ou_state, ou_new)
+        if self.framework in ["tf2", "tfe"]:
+            self.ou_state.assign_add(ou_new)
+            ou_state_new = self.ou_state
+        else:
+            ou_state_new = tf1.assign_add(self.ou_state, ou_new)
         high_m_low = self.action_space.high - self.action_space.low
         high_m_low = tf.where(
             tf.math.is_inf(high_m_low), tf.ones_like(high_m_low), high_m_low)
@@ -125,11 +129,18 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
         logp = tf.zeros(shape=(batch_size,), dtype=tf.float32)
 
         # Increment `last_timestep` by 1 (or set to `timestep`).
-        assign_op = (
-            tf1.assign_add(self.last_timestep, 1) if timestep is None else
-            tf1.assign(self.last_timestep, timestep))
-        with tf1.control_dependencies([assign_op, ou_state_new]):
+        if self.framework in ["tf2", "tfe"]:
+            if timestep is None:
+                self.last_timestep.assign_add(1)
+            else:
+                self.last_timestep = timestep
             return action, logp
+        else:
+            assign_op = (
+                tf1.assign_add(self.last_timestep, 1) if timestep is None else
+                tf1.assign(self.last_timestep, timestep))
+            with tf1.control_dependencies([assign_op, ou_state_new]):
+                return action, logp
 
     @override(GaussianNoise)
     def _get_torch_exploration_action(self, action_dist, explore, timestep):

rllib/utils/exploration/tests/test_explorations.py

@@ -27,7 +27,7 @@ def do_test_explorations(run,
         core_config["num_workers"] = 0
 
     # Test all frameworks.
-    for fw in framework_iterator(core_config):
+    for _ in framework_iterator(core_config):
         print("Agent={}".format(run))
 
         # Test for both the default Agent's exploration AND the `Random`

rllib/utils/exploration/tests/test_parameter_noise.py

@@ -189,7 +189,7 @@ class TestParameterNoise(unittest.TestCase):
 
     def _get_current_weight(self, policy, fw):
         weights = policy.get_weights()
-        key = 0 if fw == "tfe" else list(weights.keys())[0]
+        key = 0 if fw in ["tf2", "tfe"] else list(weights.keys())[0]
         return weights[key][0][0]
 
 

rllib/utils/test_utils.py

@@ -18,7 +18,7 @@ logger = logging.getLogger(__name__)
 
 
 def framework_iterator(config=None,
-                       frameworks=("tf", "tfe", "torch"),
+                       frameworks=("tf2", "tf", "tfe", "torch"),
                        session=False):
     """An generator that allows for looping through n frameworks for testing.
 
@@ -29,18 +29,23 @@ def framework_iterator(config=None,
         config (Optional[dict]): An optional config dict to alter in place
             depending on the iteration.
         frameworks (Tuple[str]): A list/tuple of the frameworks to be tested.
-            Allowed are: "tf", "tfe", "torch", and None.
+            Allowed are: "tf2", "tf", "tfe", "torch", and None.
         session (bool): If True and only in the tf-case: Enter a tf.Session()
             and yield that as second return value (otherwise yield (fw, None)).
 
     Yields:
         str: If enter_session is False:
-            The current framework ("tf", "tfe", "torch") used.
+            The current framework ("tf2", "tf", "tfe", "torch") used.
         Tuple(str, Union[None,tf.Session]: If enter_session is True:
             A tuple of the current fw and the tf.Session if fw="tf".
     """
     config = config or {}
-    frameworks = [frameworks] if isinstance(frameworks, str) else frameworks
+    frameworks = [frameworks] if isinstance(frameworks, str) else \
+        list(frameworks)
+
+    # Both tf2 and tfe present -> remove "tfe" or "tf2" depending on version.
+    if "tf2" in frameworks and "tfe" in frameworks:
+        frameworks.remove("tfe" if tfv == 2 else "tf2")
 
     for fw in frameworks:
         # Skip non-installed frameworks.
@@ -53,10 +58,14 @@ def framework_iterator(config=None,
                            "installed)!".format(fw))
             continue
         elif fw == "tfe" and not eager_mode:
-            logger.warning("framework_iterator skipping eager (could not "
+            logger.warning("framework_iterator skipping tf-eager (could not "
                            "import `eager_mode` from tensorflow.python)!")
             continue
-        assert fw in ["tf", "tfe", "torch", None]
+        elif fw == "tf2" and tfv != 2:
+            logger.warning(
+                "framework_iterator skipping tf2.x (tf version is < 2.0)!")
+            continue
+        assert fw in ["tf2", "tf", "tfe", "torch", None]
 
         # Do we need a test session?
         sess = None
@@ -69,10 +78,12 @@ def framework_iterator(config=None,
         config["framework"] = fw
 
         eager_ctx = None
-        if fw == "tfe":
+        # Enable eager mode for tf2 and tfe.
+        if fw in ["tf2", "tfe"]:
             eager_ctx = eager_mode()
             eager_ctx.__enter__()
             assert tf1.executing_eagerly()
+        # Make sure, eager mode is off.
         elif fw == "tf":
             assert not tf1.executing_eagerly()
 
@@ -169,8 +180,13 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False):
         if tf1 is not None:
             # y should never be a Tensor (y=expected value).
             if isinstance(y, tf1.Tensor):
-                raise ValueError("`y` (expected value) must not be a Tensor. "
-                                 "Use numpy.ndarray instead")
+                # In eager mode, numpyize tensors.
+                if tf.executing_eagerly():
+                    y = y.numpy()
+                else:
+                    raise ValueError(
+                        "`y` (expected value) must not be a Tensor. "
+                        "Use numpy.ndarray instead")
             if isinstance(x, tf1.Tensor):
                 # In eager mode, numpyize tensors.
                 if tf1.executing_eagerly():

rllib/utils/tf_ops.py

@@ -32,11 +32,18 @@ def minimize_and_clip(optimizer, objective, var_list, clip_val=10.0):
     # Accidentally passing values < 0.0 will break all gradients.
     assert clip_val > 0.0, clip_val
 
-    gradients = optimizer.compute_gradients(objective, var_list=var_list)
-    for i, (grad, var) in enumerate(gradients):
+    if tf.executing_eagerly():
+        tape = optimizer.tape
+        grads_and_vars = list(zip(list(
+            tape.gradient(objective, var_list)), var_list))
+    else:
+        grads_and_vars = optimizer.compute_gradients(
+            objective, var_list=var_list)
+
+    for i, (grad, var) in enumerate(grads_and_vars):
         if grad is not None:
-            gradients[i] = (tf.clip_by_norm(grad, clip_val), var)
-    return gradients
+            grads_and_vars[i] = (tf.clip_by_norm(grad, clip_val), var)
+    return grads_and_vars
 
 
 def make_tf_callable(session_or_none, dynamic_shape=False):