[RLlib] Experiment with py_func as a means to further unify tf and torch (Schedule classes). (#6951)

rllib/utils/schedules/constant_schedule.py

@@ -14,5 +14,5 @@ class ConstantSchedule(Schedule):
         super().__init__(framework=None)
         self._v = value
 
-    def value(self, t=None):
+    def _value(self, t=None):
         return self._v

rllib/utils/schedules/exponential_schedule.py

@@ -28,7 +28,7 @@ class ExponentialSchedule(Schedule):
         self.initial_p = initial_p
         self.decay_rate = decay_rate
 
-    def value(self, t):
+    def _value(self, t):
         """
         Returns the result of:
         initial_p * decay_rate ** (`t`/t_max)

rllib/utils/schedules/piecewise_schedule.py

@@ -33,9 +33,7 @@ class PiecewiseSchedule(Schedule):
                 returned. If None then an AssertionError is raised when outside
                 value is requested.
         """
-        # TODO(sven): support tf.
-        assert framework is None
-        super().__init__(framework=None)
+        super().__init__(framework=framework)
 
         idxes = [e[0] for e in endpoints]
         assert idxes == sorted(idxes)
@@ -43,7 +41,7 @@ class PiecewiseSchedule(Schedule):
         self.outside_value = outside_value
         self.endpoints = endpoints
 
-    def value(self, t):
+    def _value(self, t):
         for (l_t, l), (r_t, r) in zip(self.endpoints[:-1], self.endpoints[1:]):
             if l_t <= t < r_t:
                 alpha = float(t - l_t) / (r_t - l_t)

rllib/utils/schedules/polynomial_schedule.py

@@ -30,17 +30,10 @@ class PolynomialSchedule(Schedule):
         self.initial_p = initial_p
         self.power = power
 
-    def value(self, t):
+    def _value(self, t):
         """
         Returns the result of:
         final_p + (initial_p - final_p) * (1 - `t`/t_max) ** power
         """
-        if self.framework == "tf" and tf.executing_eagerly() is False:
-            return tf.train.polynomial_decay(
-                learning_rate=self.initial_p,
-                global_step=t,
-                decay_steps=self.schedule_timesteps,
-                end_learning_rate=self.final_p,
-                power=self.power)
         return self.final_p + (self.initial_p - self.final_p) * (
             1.0 - (t / self.schedule_timesteps))**self.power

rllib/utils/schedules/schedule.py

@@ -1,6 +1,9 @@
 from abc import ABCMeta, abstractmethod
 
 from ray.rllib.utils.framework import check_framework
+from ray.rllib.utils.framework import try_import_tf
+
+tf = try_import_tf()
 
 
 class Schedule(metaclass=ABCMeta):
@@ -26,19 +29,26 @@ class Schedule(metaclass=ABCMeta):
         self.framework = check_framework(framework)
 
     @abstractmethod
-    def value(self, t):
+    def _value(self, t):
         """
-        Returns the value based on a time value.
+        Returns the value based on a time step input.
 
         Args:
-            t (int): The time value (e.g. a time step).
-                NOTE: This could be a tf.Tensor.
+            t (int): The time step. This could be a tf.Tensor.
 
         Returns:
             any: The calculated value depending on the schedule and `t`.
         """
         raise NotImplementedError
 
+    def value(self, t):
+        if self.framework == "tf" and tf.executing_eagerly() is False:
+            return tf.cast(
+                tf.py_func(self._value, [t], tf.float64),
+                tf.float32,
+                name="schedule-value")
+        return self._value(t)
+
     def __call__(self, t):
         """
         Simply calls `self.value(t)`.

rllib/utils/schedules/tests/test_schedules.py

@@ -1,3 +1,4 @@
+from tensorflow.python.eager.context import eager_mode
 import unittest
 
 from ray.rllib.utils.schedules import ConstantSchedule, \
@@ -17,69 +18,96 @@ class TestSchedules(unittest.TestCase):
         value = 2.3
         ts = [100, 0, 10, 2, 3, 4, 99, 56, 10000, 23, 234, 56]
 
+        config = {"value": value}
+
         for fw in ["tf", "torch", None]:
-            constant = from_config(ConstantSchedule,
-                                   dict(value=value, framework=fw))
+            constant = from_config(ConstantSchedule, config, framework=fw)
+            for t in ts:
+                out = constant(t)
+                check(out, value)
+
+        # Test eager as well.
+        with eager_mode():
+            constant = from_config(ConstantSchedule, config, framework="tf")
             for t in ts:
                 out = constant(t)
                 check(out, value)
 
     def test_linear_schedule(self):
         ts = [0, 50, 10, 100, 90, 2, 1, 99, 23]
+        config = {"schedule_timesteps": 100, "initial_p": 2.1, "final_p": 0.6}
         for fw in ["tf", "torch", None]:
-            linear = from_config(
-                LinearSchedule, {
-                    "schedule_timesteps": 100,
-                    "initial_p": 2.1,
-                    "final_p": 0.6,
-                    "framework": fw
-                })
-            if fw == "tf":
-                tf.enable_eager_execution()
+            linear = from_config(LinearSchedule, config, framework=fw)
+            for t in ts:
+                out = linear(t)
+                check(out, 2.1 - (t / 100) * (2.1 - 0.6), decimals=4)
+
+        # Test eager as well.
+        with eager_mode():
+            linear = from_config(LinearSchedule, config, framework="tf")
             for t in ts:
                 out = linear(t)
                 check(out, 2.1 - (t / 100) * (2.1 - 0.6), decimals=4)
 
     def test_polynomial_schedule(self):
         ts = [0, 5, 10, 100, 90, 2, 1, 99, 23]
+        config = dict(
+            type="ray.rllib.utils.schedules.polynomial_schedule."
+            "PolynomialSchedule",
+            schedule_timesteps=100,
+            initial_p=2.0,
+            final_p=0.5,
+            power=2.0)
         for fw in ["tf", "torch", None]:
-            polynomial = from_config(
-                dict(
-                    type="ray.rllib.utils.schedules.polynomial_schedule."
-                    "PolynomialSchedule",
-                    schedule_timesteps=100,
-                    initial_p=2.0,
-                    final_p=0.5,
-                    power=2.0,
-                    framework=fw))
-            if fw == "tf":
-                tf.enable_eager_execution()
+            config["framework"] = fw
+            polynomial = from_config(config)
+            for t in ts:
+                out = polynomial(t)
+                check(out, 0.5 + (2.0 - 0.5) * (1.0 - t / 100)**2, decimals=4)
+
+        # Test eager as well.
+        with eager_mode():
+            config["framework"] = "tf"
+            polynomial = from_config(config)
             for t in ts:
                 out = polynomial(t)
                 check(out, 0.5 + (2.0 - 0.5) * (1.0 - t / 100)**2, decimals=4)
 
     def test_exponential_schedule(self):
         ts = [0, 5, 10, 100, 90, 2, 1, 99, 23]
+        config = dict(initial_p=2.0, decay_rate=0.99, schedule_timesteps=100)
         for fw in ["tf", "torch", None]:
-            exponential = from_config(
-                ExponentialSchedule,
-                dict(
-                    initial_p=2.0,
-                    decay_rate=0.99,
-                    schedule_timesteps=100,
-                    framework=fw))
+            config["framework"] = fw
+            exponential = from_config(ExponentialSchedule, config)
+            for t in ts:
+                out = exponential(t)
+                check(out, 2.0 * 0.99**(t / 100), decimals=4)
+
+        # Test eager as well.
+        with eager_mode():
+            config["framework"] = "tf"
+            exponential = from_config(ExponentialSchedule, config)
             for t in ts:
                 out = exponential(t)
                 check(out, 2.0 * 0.99**(t / 100), decimals=4)
 
     def test_piecewise_schedule(self):
-        piecewise = from_config(
-            PiecewiseSchedule,
-            dict(
-                endpoints=[(0, 50.0), (25, 100.0), (30, 200.0)],
-                outside_value=14.5))
         ts = [0, 5, 10, 100, 90, 2, 1, 99, 27]
         expected = [50.0, 60.0, 70.0, 14.5, 14.5, 54.0, 52.0, 14.5, 140.0]
-        for t, e in zip(ts, expected):
-            out = piecewise(t)
-            check(out, e, decimals=4)
+        config = dict(
+            endpoints=[(0, 50.0), (25, 100.0), (30, 200.0)],
+            outside_value=14.5)
+        for fw in ["tf", "torch", None]:
+            config["framework"] = fw
+            piecewise = from_config(PiecewiseSchedule, config)
+            for t, e in zip(ts, expected):
+                out = piecewise(t)
+                check(out, e, decimals=4)
+
+        # Test eager as well.
+        with eager_mode():
+            config["framework"] = "tf"
+            piecewise = from_config(PiecewiseSchedule, config)
+            for t, e in zip(ts, expected):
+                out = piecewise(t)
+                check(out, e, decimals=4)

rllib/utils/test_utils.py

@@ -13,8 +13,10 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False):
     after the floating point. Uses assertions.
 
     Args:
-        x (any): The first value to be compared (to `y`).
-        y (any): The second value to be compared (to `x`).
+        x (any): The value to be compared (to the expectation: `y`). This
+            may be a Tensor.
+        y (any): The expected value to be compared to `x`. This must not
+            be a Tensor.
         decimals (int): The number of digits after the floating point up to
             which all numeric values have to match.
         atol (float): Absolute tolerance of the difference between x and y
@@ -84,11 +86,26 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False):
                 raise e
     # Everything else (assume numeric).
     else:
-        # Numpyize tensors if necessary.
-        if tf is not None and isinstance(x, tf.Tensor):
-            x = x.numpy()
-        if tf is not None and isinstance(y, tf.Tensor):
-            y = y.numpy()
+        if tf is not None:
+            # y should never be a Tensor (y=expected value).
+            if isinstance(y, tf.Tensor):
+                raise ValueError("`y` (expected value) must not be a Tensor. "
+                                 "Use numpy.ndarray instead")
+            if isinstance(x, tf.Tensor):
+                # In eager mode, numpyize tensors.
+                if tf.executing_eagerly():
+                    x = x.numpy()
+                # Otherwise, ???
+                else:
+                    with tf.Session() as sess:
+                        x = sess.run(x)
+                        check(
+                            x,
+                            y,
+                            decimals=decimals,
+                            atol=atol,
+                            rtol=rtol,
+                            false=false)
 
         # Using decimals.
         if atol is None and rtol is None: