From 136ada5fb976b4474abd9d890392c135ae912757 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Thu, 30 Jan 2020 20:27:57 +0100 Subject: [PATCH] [RLlib] Experiment with py_func as a means to further unify tf and torch (Schedule classes). (#6951) --- rllib/utils/schedules/constant_schedule.py | 2 +- rllib/utils/schedules/exponential_schedule.py | 2 +- rllib/utils/schedules/piecewise_schedule.py | 6 +- rllib/utils/schedules/polynomial_schedule.py | 9 +- rllib/utils/schedules/schedule.py | 18 +++- rllib/utils/schedules/tests/test_schedules.py | 102 +++++++++++------- rllib/utils/test_utils.py | 31 ++++-- 7 files changed, 108 insertions(+), 62 deletions(-) diff --git a/rllib/utils/schedules/constant_schedule.py b/rllib/utils/schedules/constant_schedule.py index ce5d574e7..321c438d4 100644 --- a/rllib/utils/schedules/constant_schedule.py +++ b/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 diff --git a/rllib/utils/schedules/exponential_schedule.py b/rllib/utils/schedules/exponential_schedule.py index d7501b370..6cf452661 100644 --- a/rllib/utils/schedules/exponential_schedule.py +++ b/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) diff --git a/rllib/utils/schedules/piecewise_schedule.py b/rllib/utils/schedules/piecewise_schedule.py index 7ae64a62e..d03c3b0d1 100644 --- a/rllib/utils/schedules/piecewise_schedule.py +++ b/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) diff --git a/rllib/utils/schedules/polynomial_schedule.py b/rllib/utils/schedules/polynomial_schedule.py index c6aa840cf..9ddc10a44 100644 --- a/rllib/utils/schedules/polynomial_schedule.py +++ b/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 diff --git a/rllib/utils/schedules/schedule.py b/rllib/utils/schedules/schedule.py index 3bb933376..48776d274 100644 --- a/rllib/utils/schedules/schedule.py +++ b/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)`. diff --git a/rllib/utils/schedules/tests/test_schedules.py b/rllib/utils/schedules/tests/test_schedules.py index 294a2b194..5d6481202 100644 --- a/rllib/utils/schedules/tests/test_schedules.py +++ b/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) diff --git a/rllib/utils/test_utils.py b/rllib/utils/test_utils.py index 04f5d7eb5..7ea734b4f 100644 --- a/rllib/utils/test_utils.py +++ b/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: