From ae0d9bd5eb64e8da863fb29ffd182039eb8cbc51 Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Mon, 19 Jun 2017 22:31:11 -0600
Subject: [PATCH 1/6] Support batched forward and backward for FFN

---
 src/mlpack/methods/ann/ffn.hpp                | 35 ++++++---
 src/mlpack/methods/ann/ffn_impl.hpp           | 48 +++++++++---
 src/mlpack/methods/ann/layer/linear_impl.hpp  |  5 +-
 src/mlpack/tests/feedforward_network_test.cpp | 75 +++++++++++++++++++
 4 files changed, 141 insertions(+), 22 deletions(-)

diff --git a/src/mlpack/methods/ann/ffn.hpp b/src/mlpack/methods/ann/ffn.hpp
index f3f69bc4da..e64f5f8b06 100644
--- a/src/mlpack/methods/ann/ffn.hpp
+++ b/src/mlpack/methods/ann/ffn.hpp
@@ -182,16 +182,6 @@ class FFN
                 const size_t i,
                 arma::mat& gradient);
 
-  /**
-   * Compute the gradient of the feedforward network based on given input and target.
-   *
-   * @param predictors Input training variables.
-   * @param responses Outputs results from input training variables.
-   * @return Desired gradients of the feedforward network.
-   */
-  arma::mat Gradient(const arma::mat& predictors,
-                     const arma::mat& responses);
-
   /*
    * Add a new module to the model.
    *
@@ -224,6 +214,31 @@ class FFN
   template<typename Archive>
   void Serialize(Archive& ar, const unsigned int /* version */);
 
+  /**
+   * Perform the forward pass of the data in real batch mode.
+   *
+   * Forward and Backward should be used as a pair, and they are designed mainly
+   * for advanced users. User should try to use Predict and Train unless those
+   * two functions can't satisfy some special requirements.
+   *
+   * @param inputs The input data.
+   * @param results The predicted results.
+   */
+  void Forward(arma::mat inputs, arma::mat& results);
+
+  /**
+   * Perform the backward pass of the data in real batch mode.
+   *
+   * Forward and Backward should be used as a pair, and they are designed mainly
+   * for advanced users. User should try to use Predict and Train unless those
+   * two functions can't satisfy some special requirements.
+   *
+   * @param targets The training target.
+   * @param gradients Computed gradients.
+   * @return Training error of the current pass.
+   */
+  double Backward(arma::mat targets, arma::mat& gradients);
+
  private:
   // Helper functions.
   /**
diff --git a/src/mlpack/methods/ann/ffn_impl.hpp b/src/mlpack/methods/ann/ffn_impl.hpp
index 62948f1133..42de363651 100644
--- a/src/mlpack/methods/ann/ffn_impl.hpp
+++ b/src/mlpack/methods/ann/ffn_impl.hpp
@@ -130,6 +130,44 @@ void FFN<OutputLayerType, InitializationRuleType>::Train(
       << "." << std::endl;
 }
 
+template<typename OutputLayerType, typename InitializationRuleType>
+void FFN<OutputLayerType, InitializationRuleType>::Forward(
+    arma::mat inputs, arma::mat& results)
+{
+  if (parameter.is_empty())
+    ResetParameters();
+
+  if (!deterministic)
+  {
+    deterministic = true;
+    ResetDeterministic();
+  }
+
+  currentInput = std::move(inputs);
+  Forward(std::move(currentInput));
+  results = boost::apply_visitor(outputParameterVisitor, network.back());
+}
+
+template<typename OutputLayerType, typename InitializationRuleType>
+    double FFN<OutputLayerType, InitializationRuleType>::Backward(
+        arma::mat targets, arma::mat& gradients)
+{
+  currentTarget = std::move(targets);
+  double res = outputLayer.Forward(std::move(boost::apply_visitor(
+      outputParameterVisitor, network.back())), std::move(currentTarget));
+
+  outputLayer.Backward(std::move(boost::apply_visitor(outputParameterVisitor,
+      network.back())), std::move(currentTarget), std::move(error));
+
+  gradients = arma::zeros<arma::mat>(parameter.n_rows, parameter.n_cols);
+
+  Backward();
+  ResetGradients(gradients);
+  Gradient();
+
+  return res;
+}
+
 template<typename OutputLayerType, typename InitializationRuleType>
 void FFN<OutputLayerType, InitializationRuleType>::Predict(
     arma::mat predictors, arma::mat& results)
@@ -218,16 +256,6 @@ void FFN<OutputLayerType, InitializationRuleType>::Gradient(
   Gradient();
 }
 
-template<typename OutputLayerType, typename InitializationRuleType>
-arma::mat FFN<OutputLayerType, InitializationRuleType>::Gradient(
-  const arma::mat& predictors, const arma::mat& responses)
-{
-  ResetData(predictors, responses);
-  arma::mat gradients;
-  Gradient(Parameters(), 0, gradients);
-  return gradients;
-};
-
 template<typename OutputLayerType, typename InitializationRuleType>
 void FFN<OutputLayerType, InitializationRuleType>::ResetParameters()
 {
diff --git a/src/mlpack/methods/ann/layer/linear_impl.hpp b/src/mlpack/methods/ann/layer/linear_impl.hpp
index 8b726e5d25..45b5bcdefe 100644
--- a/src/mlpack/methods/ann/layer/linear_impl.hpp
+++ b/src/mlpack/methods/ann/layer/linear_impl.hpp
@@ -48,7 +48,8 @@ template<typename eT>
 void Linear<InputDataType, OutputDataType>::Forward(
     const arma::Mat<eT>&& input, arma::Mat<eT>&& output)
 {
-  output = (weight * input) + bias;
+  output = weight * input;
+  output.each_col() += bias;
 }
 
 template<typename InputDataType, typename OutputDataType>
@@ -68,7 +69,7 @@ void Linear<InputDataType, OutputDataType>::Gradient(
 {
   gradient.submat(0, 0, weight.n_elem - 1, 0) = arma::vectorise(
       error * input.t());
-  gradient.submat(weight.n_elem, 0, gradient.n_elem - 1, 0) = error;
+  gradient.submat(weight.n_elem, 0, gradient.n_elem - 1, 0) = arma::mean(error, 1);
 }
 
 template<typename InputDataType, typename OutputDataType>
diff --git a/src/mlpack/tests/feedforward_network_test.cpp b/src/mlpack/tests/feedforward_network_test.cpp
index d783830c9c..78a2dd8f0e 100644
--- a/src/mlpack/tests/feedforward_network_test.cpp
+++ b/src/mlpack/tests/feedforward_network_test.cpp
@@ -13,6 +13,7 @@
 #include <mlpack/core.hpp>
 
 #include <mlpack/core/optimizers/rmsprop/rmsprop.hpp>
+#include <mlpack/core/optimizers/sgd/update_policies/vanilla_update.hpp>
 #include <mlpack/methods/ann/layer/layer.hpp>
 #include <mlpack/methods/ann/ffn.hpp>
 
@@ -152,6 +153,80 @@ BOOST_AUTO_TEST_CASE(VanillaNetworkTest)
       (dataset, labels, dataset, labels, 2, 10, 50, 0.2);
 }
 
+BOOST_AUTO_TEST_CASE(ForwardBackwardTest)
+{
+  arma::mat dataset;
+  dataset.load("mnist_first250_training_4s_and_9s.arm");
+
+  // Normalize each point since these are images.
+  for (size_t i = 0; i < dataset.n_cols; ++i)
+    dataset.col(i) /= norm(dataset.col(i), 2);
+
+  arma::mat labels = arma::zeros(1, dataset.n_cols);
+  labels.submat(0, labels.n_cols / 2, 0, labels.n_cols - 1).fill(1);
+  labels += 1;
+
+  FFN<NegativeLogLikelihood<> > model;
+  model.Add<Linear<> >(dataset.n_rows, 50);
+  model.Add<SigmoidLayer<> >();
+  model.Add<Linear<> >(50, 10);
+  model.Add<LogSoftMax<> >();
+
+  VanillaUpdate opt;
+  model.ResetParameters();
+  opt.Initialize(model.Parameters().n_rows, model.Parameters().n_cols);
+  double stepSize = 0.01;
+  size_t batchSize = 10;
+
+  size_t iteration = 0;
+  bool converged = false;
+  while (iteration < 1000)
+  {
+    arma::running_stat<double> error;
+    size_t batchStart = 0;
+    while (batchStart < dataset.n_cols)
+    {
+      size_t batchEnd = std::min(batchStart + batchSize, (size_t)dataset.n_cols);
+      arma::mat currentData = dataset.cols(batchStart, batchEnd - 1);
+      arma::mat currentLabels = labels.cols(batchStart, batchEnd - 1);
+      arma::mat currentResuls;
+      model.Forward(currentData, currentResuls);
+      arma::mat gradients;
+      model.Backward(currentLabels, gradients);
+      opt.Update(model.Parameters(), stepSize, gradients);
+      batchStart = batchEnd;
+
+      arma::mat prediction = arma::zeros<arma::mat>(1, currentResuls.n_cols);
+
+      for (size_t i = 0; i < currentResuls.n_cols; ++i)
+      {
+        prediction(i) = arma::as_scalar(arma::find(
+            arma::max(currentResuls.col(i)) == currentResuls.col(i), 1)) + 1;
+      }
+
+      size_t correct = 0;
+      for (size_t i = 0; i < currentLabels.n_cols; i++)
+      {
+        if (int(arma::as_scalar(prediction.col(i))) ==
+            int(arma::as_scalar(currentLabels.col(i))))
+        {
+          correct++;
+        }
+      }
+
+      error(1 - (double)correct / batchSize);
+    }
+    Log::Debug << "Current training error: " << error.mean() << std::endl;
+    iteration++;
+    if (error.mean() < 0.01)
+    {
+      converged = true;
+      break;
+    }
+  }
+  BOOST_REQUIRE(converged);
+}
+
 /**
  * Train and evaluate a Dropout network with the specified structure.
  */

From 5ef19278c61d7758f7cc99d867469a57a71b1444 Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Mon, 19 Jun 2017 22:47:55 -0600
Subject: [PATCH 2/6] Minor style fixes

---
 src/mlpack/methods/ann/ffn_impl.hpp           | 4 ++--
 src/mlpack/methods/ann/layer/linear_impl.hpp  | 3 ++-
 src/mlpack/tests/feedforward_network_test.cpp | 3 ++-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/mlpack/methods/ann/ffn_impl.hpp b/src/mlpack/methods/ann/ffn_impl.hpp
index 42de363651..a522eff2d2 100644
--- a/src/mlpack/methods/ann/ffn_impl.hpp
+++ b/src/mlpack/methods/ann/ffn_impl.hpp
@@ -149,8 +149,8 @@ void FFN<OutputLayerType, InitializationRuleType>::Forward(
 }
 
 template<typename OutputLayerType, typename InitializationRuleType>
-    double FFN<OutputLayerType, InitializationRuleType>::Backward(
-        arma::mat targets, arma::mat& gradients)
+double FFN<OutputLayerType, InitializationRuleType>::Backward(
+    arma::mat targets, arma::mat& gradients)
 {
   currentTarget = std::move(targets);
   double res = outputLayer.Forward(std::move(boost::apply_visitor(
diff --git a/src/mlpack/methods/ann/layer/linear_impl.hpp b/src/mlpack/methods/ann/layer/linear_impl.hpp
index 45b5bcdefe..9173f1c80a 100644
--- a/src/mlpack/methods/ann/layer/linear_impl.hpp
+++ b/src/mlpack/methods/ann/layer/linear_impl.hpp
@@ -69,7 +69,8 @@ void Linear<InputDataType, OutputDataType>::Gradient(
 {
   gradient.submat(0, 0, weight.n_elem - 1, 0) = arma::vectorise(
       error * input.t());
-  gradient.submat(weight.n_elem, 0, gradient.n_elem - 1, 0) = arma::mean(error, 1);
+  gradient.submat(weight.n_elem, 0, gradient.n_elem - 1, 0) =
+      arma::mean(error, 1);
 }
 
 template<typename InputDataType, typename OutputDataType>
diff --git a/src/mlpack/tests/feedforward_network_test.cpp b/src/mlpack/tests/feedforward_network_test.cpp
index 78a2dd8f0e..4d850ac6e0 100644
--- a/src/mlpack/tests/feedforward_network_test.cpp
+++ b/src/mlpack/tests/feedforward_network_test.cpp
@@ -186,7 +186,8 @@ BOOST_AUTO_TEST_CASE(ForwardBackwardTest)
     size_t batchStart = 0;
     while (batchStart < dataset.n_cols)
     {
-      size_t batchEnd = std::min(batchStart + batchSize, (size_t)dataset.n_cols);
+      size_t batchEnd = std::min(batchStart + batchSize,
+          (size_t)dataset.n_cols);
       arma::mat currentData = dataset.cols(batchStart, batchEnd - 1);
       arma::mat currentLabels = labels.cols(batchStart, batchEnd - 1);
       arma::mat currentResuls;

From 3aafde97c1cdf68cf778bdc982b85bc87b02f3cc Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Tue, 20 Jun 2017 22:28:32 -0600
Subject: [PATCH 3/6] Update LogSoftmax

---
 src/mlpack/methods/ann/layer/log_softmax_impl.hpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/mlpack/methods/ann/layer/log_softmax_impl.hpp b/src/mlpack/methods/ann/layer/log_softmax_impl.hpp
index 418828b9f2..1e26499629 100644
--- a/src/mlpack/methods/ann/layer/log_softmax_impl.hpp
+++ b/src/mlpack/methods/ann/layer/log_softmax_impl.hpp
@@ -57,7 +57,8 @@ void LogSoftMax<InputDataType, OutputDataType>::Forward(
     return 0.0;
   });
 
-  output = input - (maxInput + std::log(arma::accu(output)));
+  maxInput.each_row() += arma::log(arma::sum(output));
+  output = input - maxInput;
 }
 
 template<typename InputDataType, typename OutputDataType>
@@ -67,7 +68,7 @@ void LogSoftMax<InputDataType, OutputDataType>::Backward(
     arma::Mat<eT>&& gy,
     arma::Mat<eT>&& g)
 {
-  g = gy - arma::exp(input) * arma::accu(gy);
+  g = arma::exp(input) + gy;
 }
 
 template<typename InputDataType, typename OutputDataType>

From dc782b784d24d7208bcefed6b87de21d4db14678 Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Thu, 22 Jun 2017 19:59:32 -0600
Subject: [PATCH 4/6] Apply batched forward/backward to q learning

---
 .../reinforcement_learning/q_learning.hpp     | 17 +++++++----
 .../q_learning_impl.hpp                       | 30 +++++++++++--------
 src/mlpack/tests/q_learning_test.cpp          | 25 +++++++---------
 3 files changed, 40 insertions(+), 32 deletions(-)

diff --git a/src/mlpack/methods/reinforcement_learning/q_learning.hpp b/src/mlpack/methods/reinforcement_learning/q_learning.hpp
index 4f89e7b450..27fbed044d 100644
--- a/src/mlpack/methods/reinforcement_learning/q_learning.hpp
+++ b/src/mlpack/methods/reinforcement_learning/q_learning.hpp
@@ -42,14 +42,14 @@ namespace rl {
  *
  * @tparam EnvironmentType The environment of the reinforcement learning task.
  * @tparam NetworkType The network to compute action value.
- * @tparam OptimizerType The optimizer to train the network.
+ * @tparam UpdaterType How to apply gradients when training.
  * @tparam PolicyType Behavior policy of the agent.
  * @tparam ReplayType Experience replay method.
  */
 template <
   typename EnvironmentType,
   typename NetworkType,
-  typename OptimizerType,
+  typename UpdaterType,
   typename PolicyType,
   typename ReplayType = RandomReplay<EnvironmentType>
 >
@@ -69,7 +69,7 @@ class QLearning
    * object, be sure to use std::move to avoid unnecessary copy.
    *
    * @param network The network to compute action value.
-   * @param optimizer The optimizer to train the network.
+   * @param stepSize Learning rate.
    * @param discount Discount for future return.
    * @param policy Behavior policy of the agent.
    * @param replayMethod Experience replay method.
@@ -78,10 +78,11 @@ class QLearning
    * @param explorationSteps Steps before starting to learn.
    * @param doubleQLearning Whether to use double Q-Learning.
    * @param stepLimit Maximum steps in each episode, 0 means no limit.
+   * @param updater How to apply gradients when training.
    * @param environment Reinforcement learning task.
    */
   QLearning(NetworkType network,
-            OptimizerType optimizer,
+            const double stepSize,
             const double discount,
             PolicyType policy,
             ReplayType replayMethod,
@@ -89,6 +90,7 @@ class QLearning
             const size_t explorationSteps,
             const bool doubleQLearning = false,
             const size_t stepLimit = 0,
+            UpdaterType updater = UpdaterType(),
             EnvironmentType environment = EnvironmentType());
 
   /**
@@ -127,8 +129,11 @@ class QLearning
   //! Locally-stored target network.
   NetworkType targetNetwork;
 
-  //! Locally-stored optimizer.
-  OptimizerType optimizer;
+  //! Locally-stored learning rate.
+  double stepSize;
+
+  //! Locally-stored updater.
+  UpdaterType updater;
 
   //! Discount factor of future return.
   double discount;
diff --git a/src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp b/src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp
index f6a340ee0e..226c13b1cb 100644
--- a/src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp
+++ b/src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp
@@ -20,18 +20,18 @@ namespace rl {
 template <
   typename EnvironmentType,
   typename NetworkType,
-  typename OptimizerType,
+  typename UpdaterType,
   typename PolicyType,
   typename ReplayType
 >
 QLearning<
   EnvironmentType,
   NetworkType,
-  OptimizerType,
+  UpdaterType,
   PolicyType,
   ReplayType
 >::QLearning(NetworkType network,
-             OptimizerType optimizer,
+             const double stepSize,
              const double discount,
              PolicyType policy,
              ReplayType replayMethod,
@@ -39,9 +39,11 @@ QLearning<
              const size_t explorationsSteps,
              const bool doubleQLearning,
              const size_t stepLimit,
+             UpdaterType updater,
              EnvironmentType environment):
     learningNetwork(std::move(network)),
-    optimizer(std::move(optimizer)),
+    stepSize(stepSize),
+    updater(std::move(updater)),
     discount(discount),
     policy(std::move(policy)),
     replayMethod(std::move(replayMethod)),
@@ -54,20 +56,22 @@ QLearning<
     deterministic(false)
 {
   learningNetwork.ResetParameters();
+  this->updater.Initialize(learningNetwork.Parameters().n_rows,
+      learningNetwork.Parameters().n_cols);
   targetNetwork = learningNetwork;
 }
 
 template <
   typename EnvironmentType,
   typename NetworkType,
-  typename OptimizerType,
+  typename UpdaterType,
   typename PolicyType,
   typename ReplayType
 >
 arma::Col<size_t> QLearning<
   EnvironmentType,
   NetworkType,
-  OptimizerType,
+  UpdaterType,
   PolicyType,
   ReplayType
 >::BestAction(const arma::mat& actionValues)
@@ -85,14 +89,14 @@ arma::Col<size_t> QLearning<
 template <
   typename EnvironmentType,
   typename NetworkType,
-  typename OptimizerType,
+  typename UpdaterType,
   typename BehaviorPolicyType,
   typename ReplayType
 >
 double QLearning<
   EnvironmentType,
   NetworkType,
-  OptimizerType,
+  UpdaterType,
   BehaviorPolicyType,
   ReplayType
 >::Step()
@@ -148,7 +152,7 @@ double QLearning<
 
   // Compute the update target.
   arma::mat target;
-  learningNetwork.Predict(sampledStates, target);
+  learningNetwork.Forward(sampledStates, target);
   for (size_t i = 0; i < sampledNextStates.n_cols; ++i)
   {
     target(sampledActions[i], i) = sampledRewards[i] +
@@ -156,7 +160,9 @@ double QLearning<
   }
 
   // Learn form experience.
-  learningNetwork.Train(sampledStates, target, optimizer);
+  arma::mat gradients;
+  learningNetwork.Backward(target, gradients);
+  updater.Update(learningNetwork.Parameters(), stepSize, gradients);
 
   return reward;
 }
@@ -164,14 +170,14 @@ double QLearning<
 template <
   typename EnvironmentType,
   typename NetworkType,
-  typename OptimizerType,
+  typename UpdaterType,
   typename BehaviorPolicyType,
   typename ReplayType
 >
 double QLearning<
   EnvironmentType,
   NetworkType,
-  OptimizerType,
+  UpdaterType,
   BehaviorPolicyType,
   ReplayType
 >::Episode()
diff --git a/src/mlpack/tests/q_learning_test.cpp b/src/mlpack/tests/q_learning_test.cpp
index eede05c56b..fc19dc6909 100644
--- a/src/mlpack/tests/q_learning_test.cpp
+++ b/src/mlpack/tests/q_learning_test.cpp
@@ -12,7 +12,6 @@
 
 #include <mlpack/core.hpp>
 
-#include <mlpack/core/optimizers/gradient_descent/gradient_descent.hpp>
 #include <mlpack/methods/ann/ffn.hpp>
 #include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
 #include <mlpack/methods/ann/layer/layer.hpp>
@@ -20,6 +19,8 @@
 #include <mlpack/methods/reinforcement_learning/environment/mountain_car.hpp>
 #include <mlpack/methods/reinforcement_learning/environment/cart_pole.hpp>
 #include <mlpack/methods/reinforcement_learning/policy/greedy_policy.hpp>
+#include <mlpack/core/optimizers/adam/adam_update.hpp>
+#include <mlpack/core/optimizers/rmsprop/rmsprop_update.hpp>
 
 #include <boost/test/unit_test.hpp>
 #include "test_tools.hpp"
@@ -35,23 +36,21 @@ BOOST_AUTO_TEST_SUITE(QLearningTest);
 BOOST_AUTO_TEST_CASE(CartPoleWithDQN)
 {
   // Set up the network.
-  FFN<MeanSquaredError<>, GaussianInitialization> model;
+  FFN<MeanSquaredError<>, GaussianInitialization> model(MeanSquaredError<>(),
+      GaussianInitialization(0, 0.001));
   model.Add<Linear<>>(4, 128);
   model.Add<ReLULayer<>>();
   model.Add<Linear<>>(128, 128);
   model.Add<ReLULayer<>>();
   model.Add<Linear<>>(128, 2);
 
-  // Set up the optimizer generator.
-  StandardSGD<decltype(model)> opt(model, 0.0001, 2);
-
   // Set up the policy and replay method.
   GreedyPolicy<CartPole> policy(1.0, 1000, 0.1);
   RandomReplay<CartPole> replayMethod(10, 10000);
 
   // Set up DQN agent.
-  QLearning<CartPole, decltype(model), decltype(opt), decltype(policy)>
-      agent(std::move(model), std::move(opt), 0.9, std::move(policy),
+  QLearning<CartPole, decltype(model), AdamUpdate, decltype(policy)>
+      agent(std::move(model), 0.01, 0.9, std::move(policy),
           std::move(replayMethod), 100, 100, false, 200);
 
   arma::running_stat<double> averageReturn;
@@ -102,23 +101,21 @@ BOOST_AUTO_TEST_CASE(CartPoleWithDoubleDQN)
   for (size_t trial = 0; trial < 4; ++trial)
   {
     // Set up the network.
-    FFN<MeanSquaredError<>, GaussianInitialization> model;
+    FFN<MeanSquaredError<>, GaussianInitialization> model(MeanSquaredError<>(),
+        GaussianInitialization(0, 0.001));
     model.Add<Linear<>>(4, 20);
     model.Add<ReLULayer<>>();
     model.Add<Linear<>>(20, 20);
     model.Add<ReLULayer<>>();
     model.Add<Linear<>>(20, 2);
 
-    // Set up the optimizer.
-    StandardSGD<decltype(model)> opt(model, 0.0001, 2);
-
     // Set up the policy and replay method.
     GreedyPolicy<CartPole> policy(1.0, 1000, 0.1);
     RandomReplay<CartPole> replayMethod(10, 10000);
 
     // Set up the DQN agent.
-    QLearning<CartPole, decltype(model), decltype(opt), decltype(policy)>
-        agent(std::move(model), std::move(opt), 0.9, std::move(policy),
+    QLearning<CartPole, decltype(model), RMSPropUpdate, decltype(policy)>
+        agent(std::move(model), 0.01, 0.9, std::move(policy),
             std::move(replayMethod), 100, 100, true, 200);
 
     arma::running_stat<double> averageReturn;
@@ -134,7 +131,7 @@ BOOST_AUTO_TEST_CASE(CartPoleWithDoubleDQN)
        */
       Log::Debug << "Average return: " << averageReturn.mean()
           << " Episode return: " << episodeReturn << std::endl;
-      if (averageReturn.mean() > 30)
+      if (averageReturn.mean() > 40)
       {
         agent.Deterministic() = true;
         arma::running_stat<double> testReturn;

From e92b0bd8c2eec33061581ef45dc8ecac8c6330da Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Thu, 22 Jun 2017 20:12:08 -0600
Subject: [PATCH 5/6] Update Add

---
 src/mlpack/methods/ann/layer/add_impl.hpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/mlpack/methods/ann/layer/add_impl.hpp b/src/mlpack/methods/ann/layer/add_impl.hpp
index 3ce562007a..b04bb8247f 100644
--- a/src/mlpack/methods/ann/layer/add_impl.hpp
+++ b/src/mlpack/methods/ann/layer/add_impl.hpp
@@ -31,7 +31,8 @@ template<typename eT>
 void Add<InputDataType, OutputDataType>::Forward(
     const arma::Mat<eT>&& input, arma::Mat<eT>&& output)
 {
-  output = input + weights;
+  output = input;
+  output.each_col() += weights;
 }
 
 template<typename InputDataType, typename OutputDataType>

From b471635ade0447dc977af0542fce92cd7f7e7783 Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Thu, 22 Jun 2017 23:01:15 -0600
Subject: [PATCH 6/6] Fix failed softmax layer test

---
 src/mlpack/tests/ann_layer_test.cpp | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/mlpack/tests/ann_layer_test.cpp b/src/mlpack/tests/ann_layer_test.cpp
index 08d1cc4537..3a0e599b94 100644
--- a/src/mlpack/tests/ann_layer_test.cpp
+++ b/src/mlpack/tests/ann_layer_test.cpp
@@ -912,15 +912,18 @@ BOOST_AUTO_TEST_CASE(SimpleLogSoftmaxLayerTest)
   LogSoftMax<> module;
 
   // Test the Forward function.
-  input = arma::mat("0.5 0.5");
+  input = arma::mat("0.5; 0.5");
   module.Forward(std::move(input), std::move(output));
   BOOST_REQUIRE_SMALL(arma::accu(arma::abs(
-    arma::mat("-0.6931 -0.6931") - output)), 1e-3);
+    arma::mat("-0.6931; -0.6931") - output)), 1e-3);
 
   // Test the Backward function.
   error = arma::zeros(input.n_rows, input.n_cols);
+  // Assume LogSoftmax layer is always associated with NLL output layer.
+  error(1, 0) = -1;
   module.Backward(std::move(input), std::move(error), std::move(delta));
-  BOOST_REQUIRE_EQUAL(arma::accu(delta), 0);
+  BOOST_REQUIRE_SMALL(arma::accu(arma::abs(
+      arma::mat("1.6487; 0.6487") - delta)), 1e-3);
 }
 
 BOOST_AUTO_TEST_SUITE_END();