[tune] Implement median stopping rule (#1170)

* trial scheduler interface * remove * wip median stopping * remove * median stopping rule * update * docs * update * Revrt * update * comments * fix tesT
2026-09-10 12:38:43 +08:00 · 2017-11-03 11:25:02 -07:00
parent fdf069bd1d
commit d06beacd84
5 changed files with 225 additions and 1 deletions
@@ -0,0 +1,124 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import unittest
+
+from ray.tune.result import TrainingResult
+from ray.tune.trial import Trial
+from ray.tune.trial_scheduler import MedianStoppingRule, TrialScheduler
+
+
+def result(t, rew):
+    return TrainingResult(time_total_s=t, episode_reward_mean=rew)
+
+
+class EarlyStoppingSuite(unittest.TestCase):
+    def basicSetup(self, rule):
+        t1 = Trial("t1", "PPO")  # mean is 450, max 900, t_max=10
+        t2 = Trial("t2", "PPO")  # mean is 450, max 450, t_max=5
+        for i in range(10):
+            self.assertEqual(
+                rule.on_trial_result(None, t1, result(i, i * 100)),
+                TrialScheduler.CONTINUE)
+        for i in range(5):
+            self.assertEqual(
+                rule.on_trial_result(None, t2, result(i, 450)),
+                TrialScheduler.CONTINUE)
+        return t1, t2
+
+    def testMedianStoppingConstantPerf(self):
+        rule = MedianStoppingRule(grace_period=0, min_samples_required=1)
+        t1, t2 = self.basicSetup(rule)
+        rule.on_trial_complete(None, t1, result(10, 1000))
+        self.assertEqual(
+            rule.on_trial_result(None, t2, result(5, 450)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(
+            rule.on_trial_result(None, t2, result(6, 0)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(
+            rule.on_trial_result(None, t2, result(10, 450)),
+            TrialScheduler.STOP)
+
+    def testMedianStoppingOnCompleteOnly(self):
+        rule = MedianStoppingRule(grace_period=0, min_samples_required=1)
+        t1, t2 = self.basicSetup(rule)
+        self.assertEqual(
+            rule.on_trial_result(None, t2, result(100, 0)),
+            TrialScheduler.CONTINUE)
+        rule.on_trial_complete(None, t1, result(10, 1000))
+        self.assertEqual(
+            rule.on_trial_result(None, t2, result(101, 0)),
+            TrialScheduler.STOP)
+
+    def testMedianStoppingGracePeriod(self):
+        rule = MedianStoppingRule(grace_period=2.5, min_samples_required=1)
+        t1, t2 = self.basicSetup(rule)
+        rule.on_trial_complete(None, t1, result(10, 1000))
+        rule.on_trial_complete(None, t2, result(10, 1000))
+        t3 = Trial("t3", "PPO")
+        self.assertEqual(
+            rule.on_trial_result(None, t3, result(1, 10)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(
+            rule.on_trial_result(None, t3, result(2, 10)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(
+            rule.on_trial_result(None, t3, result(3, 10)),
+            TrialScheduler.STOP)
+
+    def testMedianStoppingMinSamples(self):
+        rule = MedianStoppingRule(grace_period=0, min_samples_required=2)
+        t1, t2 = self.basicSetup(rule)
+        rule.on_trial_complete(None, t1, result(10, 1000))
+        t3 = Trial("t3", "PPO")
+        self.assertEqual(
+            rule.on_trial_result(None, t3, result(3, 10)),
+            TrialScheduler.CONTINUE)
+        rule.on_trial_complete(None, t2, result(10, 1000))
+        self.assertEqual(
+            rule.on_trial_result(None, t3, result(3, 10)),
+            TrialScheduler.STOP)
+
+    def testMedianStoppingUsesMedian(self):
+        rule = MedianStoppingRule(grace_period=0, min_samples_required=1)
+        t1, t2 = self.basicSetup(rule)
+        rule.on_trial_complete(None, t1, result(10, 1000))
+        rule.on_trial_complete(None, t2, result(10, 1000))
+        t3 = Trial("t3", "PPO")
+        self.assertEqual(
+            rule.on_trial_result(None, t3, result(1, 260)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(
+            rule.on_trial_result(None, t3, result(2, 260)),
+            TrialScheduler.STOP)
+
+    def testAlternateMetrics(self):
+        def result2(t, rew):
+            return TrainingResult(training_iteration=t, neg_mean_loss=rew)
+
+        rule = MedianStoppingRule(
+            grace_period=0, min_samples_required=1,
+            time_attr='training_iteration', reward_attr='neg_mean_loss')
+        t1 = Trial("t1", "PPO")  # mean is 450, max 900, t_max=10
+        t2 = Trial("t2", "PPO")  # mean is 450, max 450, t_max=5
+        for i in range(10):
+            self.assertEqual(
+                rule.on_trial_result(None, t1, result2(i, i * 100)),
+                TrialScheduler.CONTINUE)
+        for i in range(5):
+            self.assertEqual(
+                rule.on_trial_result(None, t2, result2(i, 450)),
+                TrialScheduler.CONTINUE)
+        rule.on_trial_complete(None, t1, result2(10, 1000))
+        self.assertEqual(
+            rule.on_trial_result(None, t2, result2(5, 450)),
+            TrialScheduler.CONTINUE)
+        self.assertEqual(
+            rule.on_trial_result(None, t2, result2(6, 0)),
+            TrialScheduler.CONTINUE)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)