From e6cd8b621746a14a87d2ff79d02eee8924b4e9c5 Mon Sep 17 00:00:00 2001
From: Marcus Edel <marcus.edel@fu-berlin.de>
Date: Sun, 1 Sep 2019 15:32:50 +0200
Subject: [PATCH] Run the DQN reward-clipping test multiple times to account
 for minimal model size and short test time.

---
 src/mlpack/tests/q_learning_test.cpp      |   2 +-
 src/mlpack/tests/reward_clipping_test.cpp | 129 ++++++++++++----------
 2 files changed, 70 insertions(+), 61 deletions(-)

diff --git a/src/mlpack/tests/q_learning_test.cpp b/src/mlpack/tests/q_learning_test.cpp
index 02fd0e92dc..630c8e0c9d 100644
--- a/src/mlpack/tests/q_learning_test.cpp
+++ b/src/mlpack/tests/q_learning_test.cpp
@@ -105,7 +105,7 @@ BOOST_AUTO_TEST_CASE(CartPoleWithDQN)
 
 //! Test DQN in Cart Pole task with Prioritized Replay.
 BOOST_AUTO_TEST_CASE(CartPoleWithDQNPrioritizedReplay)
-  {
+{
   // Set up the network.
   FFN<MeanSquaredError<>, GaussianInitialization> model(MeanSquaredError<>(),
       GaussianInitialization(0, 0.001));
diff --git a/src/mlpack/tests/reward_clipping_test.cpp b/src/mlpack/tests/reward_clipping_test.cpp
index e219b4461e..ba05095915 100644
--- a/src/mlpack/tests/reward_clipping_test.cpp
+++ b/src/mlpack/tests/reward_clipping_test.cpp
@@ -57,72 +57,81 @@ BOOST_AUTO_TEST_CASE(ClippedRewardTest)
 //! Test DQN in Acrobot task.
 BOOST_AUTO_TEST_CASE(RewardClippedAcrobotWithDQN)
 {
-    // Set up the network.
-  FFN<MeanSquaredError<>, GaussianInitialization> model(MeanSquaredError<>(),
-      GaussianInitialization(0, 0.001));
-  model.Add<Linear<>>(4, 64);
-  model.Add<ReLULayer<>>();
-  model.Add<Linear<>>(64, 32);
-  model.Add<ReLULayer<>>();
-  model.Add<Linear<>>(32, 3);
-
-  // Set up the policy and replay method.
-  GreedyPolicy<RewardClipping<Acrobot>> policy(1.0, 1000, 0.1, 0.99);
-  RandomReplay<RewardClipping<Acrobot>> replayMethod(20, 10000);
-
-  // Set up Acrobot task and reward clipping wrapper
-  Acrobot task;
-  RewardClipping<Acrobot> rewardClipping(task, -2.0, +2.0);
-
-  // Set up update rule
-  AdamUpdate update;
-
-  TrainingConfig config;
-  config.StepSize() = 0.01;
-  config.Discount() = 0.99;
-  config.TargetNetworkSyncInterval() = 100;
-  config.ExplorationSteps() = 100;
-  config.DoubleQLearning() = false;
-  config.StepLimit() = 400;
-
-  // Set up DQN agent.
-  QLearning<decltype(rewardClipping), decltype(model), AdamUpdate,
-            decltype(policy)>
-      agent(std::move(config), std::move(model), std::move(policy),
-      std::move(replayMethod), std::move(update), std::move(rewardClipping));
-
-  arma::running_stat<double> averageReturn;
-  size_t episodes = 0;
-  bool converged = true;
-  while (true)
+  // We will allow three trials, although it would be very uncommon for the test
+  // to use more than one.
+  bool converged = false;
+  for (size_t trial = 0; trial < 3; ++trial)
   {
-    double episodeReturn = agent.Episode();
-    averageReturn(episodeReturn);
-    episodes += 1;
+    // Set up the network.
+    FFN<MeanSquaredError<>, GaussianInitialization> model(MeanSquaredError<>(),
+        GaussianInitialization(0, 0.001));
+    model.Add<Linear<>>(4, 64);
+    model.Add<ReLULayer<>>();
+    model.Add<Linear<>>(64, 32);
+    model.Add<ReLULayer<>>();
+    model.Add<Linear<>>(32, 3);
 
-    if (episodes > 1000)
+    // Set up the policy and replay method.
+    GreedyPolicy<RewardClipping<Acrobot>> policy(1.0, 1000, 0.1, 0.99);
+    RandomReplay<RewardClipping<Acrobot>> replayMethod(20, 10000);
+
+    // Set up Acrobot task and reward clipping wrapper.
+    Acrobot task;
+    RewardClipping<Acrobot> rewardClipping(task, -2.0, +2.0);
+
+    // Set up update rule.
+    AdamUpdate update;
+
+    TrainingConfig config;
+    config.StepSize() = 0.01;
+    config.Discount() = 0.99;
+    config.TargetNetworkSyncInterval() = 100;
+    config.ExplorationSteps() = 100;
+    config.DoubleQLearning() = false;
+    config.StepLimit() = 400;
+
+    // Set up DQN agent.
+    QLearning<decltype(rewardClipping), decltype(model), AdamUpdate,
+              decltype(policy)>
+        agent(std::move(config), std::move(model), std::move(policy),
+        std::move(replayMethod), std::move(update), std::move(rewardClipping));
+
+    arma::running_stat<double> averageReturn;
+    size_t episodes = 0;
+    converged = true;
+    while (true)
     {
-      Log::Debug << "Acrobot with DQN failed." << std::endl;
-      converged = false;
-      break;
+      double episodeReturn = agent.Episode();
+      averageReturn(episodeReturn);
+      episodes += 1;
+
+      if (episodes > 1000)
+      {
+        Log::Debug << "Acrobot with DQN failed." << std::endl;
+        converged = false;
+        break;
+      }
+
+      /**
+       * I am using a thresold of -380 to check convegence.
+       */
+      Log::Debug << "Average return: " << averageReturn.mean()
+          << " Episode return: " << episodeReturn << std::endl;
+      if (averageReturn.mean() > -380.00)
+      {
+        agent.Deterministic() = true;
+        arma::running_stat<double> testReturn;
+        for (size_t i = 0; i < 20; ++i)
+          testReturn(agent.Episode());
+
+        Log::Debug << "Average return in deterministic test: "
+            << testReturn.mean() << std::endl;
+        break;
+      }
     }
 
-    /**
-     * I am using a thresold of -380 to check convegence.
-     */
-    Log::Debug << "Average return: " << averageReturn.mean()
-        << " Episode return: " << episodeReturn << std::endl;
-    if (averageReturn.mean() > -380.00)
-    {
-      agent.Deterministic() = true;
-      arma::running_stat<double> testReturn;
-      for (size_t i = 0; i < 20; ++i)
-        testReturn(agent.Episode());
-
-      Log::Debug << "Average return in deterministic test: "
-          << testReturn.mean() << std::endl;
+    if (converged)
       break;
-    }
   }
 
   BOOST_REQUIRE(converged);