Merge branch 'master' into 'mc_kde_error_bounds'

CODE_OF_CONDUCT.md

@@ -0,0 +1,96 @@
+# mlpack Code of Conduct
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and
+expression, level of experience, education, socio-economic status, nationality,
+personal appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+  advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies within all project spaces, and it also applies when
+an individual is representing the project or its community in public spaces.
+Examples of representing a project or community include using an official
+project e-mail address, posting via an official social media account, or acting
+as an appointed representative at an online or offline event. Representation of
+a project may be further defined and clarified by project maintainers.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at conduct@mlpack.org. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Reporting
+
+If you believe someone is violating the code of conduct we ask that you report
+it by emailing conduct@mlpack.org. All reports will be kept confidential. In
+some cases we may determine that a public statement will need to be made. If
+that's the case, the identities of all victims and reporters will remain
+confidential unless those individuals instruct us otherwise.
+
+If you are unsure whether the incident is a violation, or whether the space
+where it happened is covered by this Code of Conduct, we encourage you to still
+report it. We would much rather have a few extra reports where we decide to take
+no action, rather than miss a report of an actual violation. We do not look
+negatively on you if we find the incident is not a violation. And knowing about
+incidents that are not violations, or happen outside our spaces, can also help
+us to improve the Code of Conduct or the processes surrounding it.
+
+In your report please include:
+
+* Your contact info (so we can get in touch with you if we need to follow up)
+* Names (real, nicknames, or pseudonyms) of any individuals involved. If there
+  were other witnesses besides you, please try to include them as well.
+* When and where the incident occurred. Please be as specific as possible.
+* Your account of what occurred. If there is a publicly available record
+  (e.g. a mailing list archive or a public IRC logger) please include a link.
+* Any extra context you believe existed for the incident.
+* If you believe this incident is ongoing.
+* Any other information you believe we should have.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 1.4, available at
+https://www.contributor-covenant.org/version/1/4/code-of-conduct.html, and
+includes some aspects of the Drupal Code of Conduct.

CONTRIBUTING.md

@@ -5,8 +5,8 @@ contribute to mlpack and join the community!  If you would like to make
 improvements to the library, add new features that are useful to you and others,
 or have found a bug that you know how to fix, please submit a pull request!
 
-If you would like to learn more about how to get started contributing, see
-[Getting Involved](http://www.mlpack.org/involved.html), and if you are
+If you would like to learn more about how to get started contributing, see the
+[Community](http://www.mlpack.org/community.html) page, and if you are
 interested in participating in Google Summer of Code, see
 [mlpack and Google Summer of Code](http://www.mlpack.org/gsoc.html).
 

COPYRIGHT.txt

@@ -113,6 +113,7 @@ Copyright:
   Copyright 2019, Dan Timson
   Copyright 2019, Miguel Canteras <mcanteras@gmail.com>
   Copyright 2019, Bishwa Karki <karkeebishwa1@gmail.com>
+  Copyright 2019, Yashwant Singh Parihar <yashwantsingh.sngh@gmail.com>
 
 License: BSD-3-clause
   All rights reserved.

Doxyfile

@@ -4,7 +4,7 @@
 # Project related configuration options
 #---------------------------------------------------------------------------
 PROJECT_NAME           = mlpack
-PROJECT_NUMBER         = 3.1.0
+PROJECT_NUMBER         = 3.1.1
 OUTPUT_DIRECTORY       = ./doc
 CREATE_SUBDIRS         = NO
 OUTPUT_LANGUAGE        = English

HISTORY.md

@@ -1,5 +1,20 @@
-### mlpack 3.1.2
+### mlpack 4.0.0
 ###### ????-??-??
+  * Add Multiple Pole Balancing Environment (#1901).
+
+  * Add new parameter `maximum_depth` to decision tree and random forest
+    bindings (#1916).
+
+  * Fix prediction output of softmax regression when test set accuracy is
+    calculated (#1922).
+
+  * Action struct in continuous RL environments now stores the action as a 
+    `double` instead of `double[1]` (#1941, #1931).
+
+  * Pendulum environment now checks for termination. All RL environments now
+    have an option to terminate after a set number of time steps (no limit
+    by default) (#1941).
+
   * Add support for probabilistic KDE (kernel density estimation) error bounds
     when using the Gaussian kernel (#1934).
 

doc/guide/python_quickstart.hpp

@@ -17,7 +17,7 @@ Installing the mlpack bindings for Python is straightforward.  It's easy to use
 conda or pip to do this:
 
 @code{.sh}
-pip install mlpack/mlpack3
+pip install mlpack3
 @endcode
 
 @code{.sh}

src/mlpack/bindings/python/setup.py.in

@@ -108,5 +108,7 @@ setup(name='mlpack',
       cmdclass={ 'build_ext': build_ext },
       ext_modules = modules,
       setup_requires=['cython', 'pytest-runner'],
-      tests_require=['pytest'],
+      tests_require=['pytest>3;python_version>"3.4"', 'pytest>3,<=4.6;python_version<="3.4"',
+                     'more-itertools>=4.0.0,<6.0.0;python_version<="2.7"',
+                     'more-itertools>=4.0.0;python_version>"2.7"'],
       zip_safe = False)

src/mlpack/methods/ann/gan/gan.hpp

@@ -210,7 +210,7 @@ class GAN
    * Gradient function for Standard GAN and DCGAN.
    * This function passes the gradient based on which network is being
    * trained, i.e., Generator or Discriminator.
-   * 
+   *
    * @param parameters present parameters of the network.
    * @param i Index of the predictors.
    * @param gradient Variable to store the present gradient.
@@ -228,7 +228,7 @@ class GAN
    * Gradient function for WGAN.
    * This function passes the gradient based on which network is being
    * trained, i.e., Generator or Discriminator.
-   * 
+   *
    * @param parameters present parameters of the network.
    * @param i Index of the predictors.
    * @param gradient Variable to store the present gradient.
@@ -245,7 +245,7 @@ class GAN
    * Gradient function for WGAN-GP.
    * This function passes the gradient based on which network is being
    * trained, i.e., Generator or Discriminator.
-   * 
+   *
    * @param parameters present parameters of the network.
    * @param i Index of the predictors.
    * @param gradient Variable to store the present gradient.
@@ -298,6 +298,16 @@ class GAN
   //! Return the number of separable functions (the number of predictor points).
   size_t NumFunctions() const { return numFunctions; }
 
+  //! Get the matrix of responses to the input data points.
+  const arma::mat& Responses() const { return responses; }
+  //! Modify the matrix of responses to the input data points.
+  arma::mat& Responses() { return responses; }
+
+  //! Get the matrix of data points (predictors).
+  const arma::mat& Predictors() const { return predictors; }
+  //! Modify the matrix of data points (predictors).
+  arma::mat& Predictors() { return predictors; }
+
   //! Serialize the model.
   template<typename Archive>
   void serialize(Archive& ar, const unsigned int /* version */);

src/mlpack/methods/ann/gan/gan_impl.hpp

@@ -42,7 +42,6 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
     const double multiplier,
     const double clippingParameter,
     const double lambda):
-    predictors(predictors),
     generator(std::move(generator)),
     discriminator(std::move(discriminator)),
     initializeRule(initializeRule),
@@ -66,17 +65,16 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
 
   this->discriminator.deterministic = this->generator.deterministic = true;
 
-  responses.set_size(1, predictors.n_cols);
-  responses.ones();
+  this->predictors.set_size(predictors.n_rows, predictors.n_cols + batchSize);
+  this->predictors.cols(0, predictors.n_cols - 1) = predictors;
+  this->discriminator.predictors = arma::mat(this->predictors.memptr(),
+      this->predictors.n_rows, this->predictors.n_cols, false, false);
 
-  this->discriminator.predictors.set_size(predictors.n_rows,
-      predictors.n_cols + batchSize);
-  this->discriminator.predictors.cols(0, predictors.n_cols - 1) = predictors;
-
-  this->discriminator.responses.set_size(1, predictors.n_cols + batchSize);
-  this->discriminator.responses.ones();
-  this->discriminator.responses.cols(predictors.n_cols,
+  responses.ones(1, predictors.n_cols + batchSize);
+  responses.cols(predictors.n_cols,
       predictors.n_cols + batchSize - 1) = arma::zeros(1, batchSize);
+  this->discriminator.responses = arma::mat(this->responses.memptr(),
+      this->responses.n_rows, this->responses.n_cols, false, false);
 
   numFunctions = predictors.n_cols;
 
@@ -232,14 +230,14 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
   noise.imbue( [&]() { return noiseFunction();} );
   generator.Forward(std::move(noise));
 
-  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
       boost::apply_visitor(outputParameterVisitor, generator.network.back());
-  discriminator.Forward(std::move(discriminator.predictors.cols(numFunctions,
+  discriminator.Forward(std::move(predictors.cols(numFunctions,
       numFunctions + batchSize - 1)));
-  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+  responses.cols(numFunctions, numFunctions + batchSize - 1) =
       arma::zeros(1, batchSize);
 
-  currentTarget = arma::mat(discriminator.responses.memptr() + numFunctions,
+  currentTarget = arma::mat(responses.memptr() + numFunctions,
       1, batchSize, false, false);
   res += discriminator.outputLayer.Forward(
       std::move(boost::apply_visitor(
@@ -299,9 +297,9 @@ EvaluateWithGradient(const arma::mat& /* parameters */,
 
   noise.imbue( [&]() { return noiseFunction();} );
   generator.Forward(std::move(noise));
-  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
       boost::apply_visitor(outputParameterVisitor, generator.network.back());
-  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+  responses.cols(numFunctions, numFunctions + batchSize - 1) =
       arma::zeros(1, batchSize);
 
   // Get the gradients of the Generator.
@@ -313,7 +311,7 @@ EvaluateWithGradient(const arma::mat& /* parameters */,
   {
     // Minimize -log(D(G(noise))).
     // Pass the error from Discriminator to Generator.
-    discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+    responses.cols(numFunctions, numFunctions + batchSize - 1) =
         arma::ones(1, batchSize);
     discriminator.Gradient(discriminator.parameter, numFunctions,
         noiseGradientDiscriminator, batchSize);
@@ -370,7 +368,9 @@ template<
 >
 void GAN<Model, InitializationRuleType, Noise, PolicyType>::Shuffle()
 {
-  math::ShuffleData(predictors, responses, predictors, responses);
+  const arma::uvec ordering = arma::shuffle(arma::linspace<arma::uvec>(0,
+      numFunctions - 1, numFunctions));
+  predictors.cols(0, numFunctions - 1) = predictors.cols(ordering);
 }
 
 template<

src/mlpack/methods/ann/gan/wgan_impl.hpp

@@ -51,14 +51,14 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
   noise.imbue( [&]() { return noiseFunction();} );
   generator.Forward(std::move(noise));
 
-  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
       boost::apply_visitor(outputParameterVisitor, generator.network.back());
-  discriminator.Forward(std::move(discriminator.predictors.cols(numFunctions,
+  discriminator.Forward(std::move(predictors.cols(numFunctions,
       numFunctions + batchSize - 1)));
-  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+  responses.cols(numFunctions, numFunctions + batchSize - 1) =
       -arma::ones(1, batchSize);
 
-  currentTarget = arma::mat(discriminator.responses.memptr() + numFunctions,
+  currentTarget = arma::mat(responses.memptr() + numFunctions,
       1, batchSize, false, false);
   res += discriminator.outputLayer.Forward(
       std::move(boost::apply_visitor(
@@ -117,9 +117,9 @@ EvaluateWithGradient(const arma::mat& /* parameters */,
 
   noise.imbue( [&]() { return noiseFunction();} );
   generator.Forward(std::move(noise));
-  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
       boost::apply_visitor(outputParameterVisitor, generator.network.back());
-  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+  responses.cols(numFunctions, numFunctions + batchSize - 1) =
       -arma::ones(1, batchSize);
 
   // Get the gradients of the Generator.
@@ -133,7 +133,7 @@ EvaluateWithGradient(const arma::mat& /* parameters */,
   {
     // Minimize -D(G(noise)).
     // Pass the error from Discriminator to Generator.
-    discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+    responses.cols(numFunctions, numFunctions + batchSize - 1) =
         arma::ones(1, batchSize);
     discriminator.Gradient(discriminator.parameter, numFunctions,
         noiseGradientDiscriminator, batchSize);

src/mlpack/methods/ann/gan/wgangp_impl.hpp

@@ -54,14 +54,14 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
 
   arma::mat generatedData = boost::apply_visitor(outputParameterVisitor,
       generator.network.back());
-  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
       generatedData;
-  discriminator.Forward(std::move(discriminator.predictors.cols(numFunctions,
+  discriminator.Forward(std::move(predictors.cols(numFunctions,
       numFunctions + batchSize - 1)));
-  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+  responses.cols(numFunctions, numFunctions + batchSize - 1) =
       -arma::ones(1, batchSize);
 
-  currentTarget = arma::mat(discriminator.responses.memptr() + numFunctions,
+  currentTarget = arma::mat(responses.memptr() + numFunctions,
       1, batchSize, false, false);
   res += discriminator.outputLayer.Forward(
       std::move(boost::apply_visitor(
@@ -70,9 +70,9 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
 
   // Gradient Penalty is calculated here.
   double epsilon = math::Random();
-  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
       (epsilon * currentInput) + ((1.0 - epsilon) * generatedData);
-  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+  responses.cols(numFunctions, numFunctions + batchSize - 1) =
       -arma::ones(1, batchSize);
   discriminator.Gradient(discriminator.parameter, numFunctions,
       normGradientDiscriminator, batchSize);
@@ -139,15 +139,15 @@ EvaluateWithGradient(const arma::mat& /* parameters */,
 
   // Gradient Penalty is calculated here.
   double epsilon = math::Random();
-  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
       (epsilon * currentInput) + ((1.0 - epsilon) * generatedData);
-  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+  responses.cols(numFunctions, numFunctions + batchSize - 1) =
       -arma::ones(1, batchSize);
   discriminator.Gradient(discriminator.parameter, numFunctions,
       normGradientDiscriminator, batchSize);
   res += lambda * std::pow(arma::norm(normGradientDiscriminator, 2) - 1, 2);
 
-  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+  predictors.cols(numFunctions, numFunctions + batchSize - 1) =
       generatedData;
   res += discriminator.EvaluateWithGradient(discriminator.parameter,
       numFunctions, noiseGradientDiscriminator, batchSize);
@@ -157,7 +157,7 @@ EvaluateWithGradient(const arma::mat& /* parameters */,
   {
     // Minimize -D(G(noise)).
     // Pass the error from Discriminator to Generator.
-    discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+    responses.cols(numFunctions, numFunctions + batchSize - 1) =
         arma::ones(1, batchSize);
     discriminator.Gradient(discriminator.parameter, numFunctions,
         noiseGradientDiscriminator, batchSize);

src/mlpack/methods/ann/layer/CMakeLists.txt

@@ -40,6 +40,8 @@ set(SOURCES
   gru_impl.hpp
   hard_tanh.hpp
   hard_tanh_impl.hpp
+  highway.hpp
+  highway_impl.hpp
   join.hpp
   join_impl.hpp
   layer.hpp

src/mlpack/methods/ann/layer/alpha_dropout.hpp

@@ -37,7 +37,6 @@ namespace ann /** Artificial Neural Network. */ {
  *   journal = {Advances in Neural Information Processing Systems},
  *   year    = {2017}
  * }
- * }
  * @endcode
  *
  * @tparam InputDataType Type of the input data (arma::colvec, arma::mat,

src/mlpack/methods/ann/layer/batch_norm.hpp

@@ -195,6 +195,9 @@ class BatchNorm
 
   //! Locally-stored normalized input.
   OutputDataType normalized;
+
+  //! Locally-stored zero mean input.
+  OutputDataType inputMean;
 }; // class BatchNorm
 
 } // namespace ann

src/mlpack/methods/ann/layer/batch_norm_impl.hpp

@@ -22,6 +22,7 @@ namespace ann { /** Artificial Neural Network. */
 
 template<typename InputDataType, typename OutputDataType>
 BatchNorm<InputDataType, OutputDataType>::BatchNorm() :
+    size(0),
     eps(1e-8),
     loading(false),
     deterministic(false),
@@ -80,6 +81,7 @@ void BatchNorm<InputDataType, OutputDataType>::Forward(
 
     // Normalize the input.
     output = input.each_col() - mean;
+    inputMean = output;
     output.each_col() /= arma::sqrt(variance + eps);
 
     // Use Welford method to compute the sample variance and mean.
@@ -87,9 +89,9 @@ void BatchNorm<InputDataType, OutputDataType>::Forward(
     {
       count += 1;
 
-      OutputDataType delta = input.col(i) - runningMean;
-      runningMean = runningMean + delta / count;
-      runningVariance += delta % (input.col(i) - runningMean);
+      OutputDataType diff = input.col(i) - runningMean;
+      runningMean = runningMean + diff / count;
+      runningVariance += diff % (input.col(i) - runningMean);
     }
 
     // Reused in the backward and gradient step.
@@ -106,7 +108,6 @@ template<typename eT>
 void BatchNorm<InputDataType, OutputDataType>::Backward(
     const arma::Mat<eT>&& input, arma::Mat<eT>&& gy, arma::Mat<eT>&& g)
 {
-  const arma::mat inputMean = input.each_col() - mean;
   const arma::mat stdInv = 1.0 / arma::sqrt(variance + eps);
 
   // Step 1: dl / dxhat
@@ -123,8 +124,7 @@ void BatchNorm<InputDataType, OutputDataType>::Backward(
 
   // Step 3: sum (dl / dxhat * -1 / stdInv) + variance *
   // (sum -2 * (x - mu)) / m.
-  g.each_col() += (arma::sum(norm.each_col() % -stdInv, 1) + (var %
-      arma::mean(-2 * inputMean, 1))) / input.n_cols;
+  g.each_col() += arma::sum(norm.each_col() % -stdInv, 1) / input.n_cols;
 }
 
 template<typename InputDataType, typename OutputDataType>

src/mlpack/methods/ann/layer/glimpse_impl.hpp

@@ -224,6 +224,9 @@ void Glimpse<InputDataType, OutputDataType>::serialize(
   ar & BOOST_SERIALIZATION_NVP(depth);
   ar & BOOST_SERIALIZATION_NVP(scale);
   ar & BOOST_SERIALIZATION_NVP(inputWidth);
+  ar & BOOST_SERIALIZATION_NVP(inputHeight);
+  ar & BOOST_SERIALIZATION_NVP(outputWidth);
+  ar & BOOST_SERIALIZATION_NVP(outputHeight);
   ar & BOOST_SERIALIZATION_NVP(location);
 }
 

src/mlpack/methods/ann/layer/highway.hpp

@@ -0,0 +1,257 @@
+/**
+ * @file highway.hpp
+ * @author Konstantin Sidorov
+ * @author Saksham Bansal
+ *
+ * Definition of the Highway layer.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_METHODS_ANN_LAYER_HIGHWAY_HPP
+#define MLPACK_METHODS_ANN_LAYER_HIGHWAY_HPP
+
+#include <mlpack/prereqs.hpp>
+
+#include <boost/ptr_container/ptr_vector.hpp>
+
+#include "../visitor/delete_visitor.hpp"
+#include "../visitor/delta_visitor.hpp"
+#include "../visitor/output_height_visitor.hpp"
+#include "../visitor/output_parameter_visitor.hpp"
+#include "../visitor/output_width_visitor.hpp"
+
+#include "layer_types.hpp"
+#include "add_merge.hpp"
+
+namespace mlpack {
+namespace ann /** Artificial Neural Network. */ {
+
+/**
+ * Implementation of the Highway layer. The Highway class can vary its behavior
+ * between that of feed-forward fully connected network container and that
+ * of a layer which simply passes its inputs through depending on the transform
+ * gate. Note that the size of the input and output matrices of this class
+ * should be equal.
+ *
+ * For more information, refer the following paper.
+ *
+ * @code
+ * @article{Srivastava2015,
+ *   author  = {Rupesh Kumar Srivastava, Klaus Greff, Jurgen Schmidhuber},
+ *   title   = {Training Very Deep Networks},
+ *   journal = {Advances in Neural Information Processing Systems},
+ *   year    = {2015},
+ *   url     = {https://arxiv.org/abs/1507.06228},
+ * }
+ * @endcode
+ *
+ * @tparam InputDataType Type of the input data (arma::colvec, arma::mat,
+ *         arma::sp_mat or arma::cube).
+ * @tparam OutputDataType Type of the output data (arma::colvec, arma::mat,
+ *         arma::sp_mat or arma::cube).
+ */
+template <
+    typename InputDataType = arma::mat,
+    typename OutputDataType = arma::mat,
+    typename... CustomLayers>
+class Highway
+{
+ public:
+  //! Create the Highway object.
+  Highway();
+
+  /**
+   * Create the Highway object.
+   *
+   * @param inSize The number of input units.
+   * @param model Expose all the network modules.
+   */
+  Highway(const size_t inSize, const bool model = true);
+
+  //! Destroy the Highway object.
+  ~Highway();
+
+  /**
+   * Destroy all the modules added to the Highway object.
+   */
+  void DeleteModules();
+
+  /**
+   * Reset the layer parameter.
+   */
+  void Reset();
+
+  /**
+   * Ordinary feed-forward pass of a neural network, evaluating the function
+   * f(x) by propagating the activity forward through f.
+   *
+   * @param input Input data used for evaluating the specified function.
+   * @param output Resulting output activation.
+   */
+  template<typename eT>
+  void Forward(arma::Mat<eT>&& input, arma::Mat<eT>&& output);
+
+  /**
+   * Ordinary feed-backward pass of a neural network, calculating the function
+   * f(x) by propagating x backwards through f. Using the results from the
+   * feed-forward pass.
+   *
+   * @param input The propagated input activation.
+   * @param gy The backpropagated error.
+   * @param g The calculated gradient.
+   */
+  template<typename eT>
+  void Backward(const arma::Mat<eT>&& /* input */,
+                arma::Mat<eT>&& gy,
+                arma::Mat<eT>&& g);
+
+  /**
+   * Calculate the gradient using the output delta and the input activation.
+   *
+   * @param input The input parameter used for calculating the gradient.
+   * @param error The calculated error.
+   * @param gradient The calculated gradient.
+   */
+  template<typename eT>
+  void Gradient(arma::Mat<eT>&& input,
+                arma::Mat<eT>&& error,
+                arma::Mat<eT>&& gradient);
+
+  /**
+   * Add a new module to the model.
+   *
+   * @param args The layer parameter.
+   */
+  template <class LayerType, class... Args>
+  void Add(Args... args) { network.push_back(new LayerType(args...)); }
+
+  /**
+   * Add a new module to the model.
+   *
+   * @param layer The Layer to be added to the model.
+   */
+  void Add(LayerTypes<CustomLayers...> layer) { network.push_back(layer); }
+
+  //! Return the modules of the model.
+  std::vector<LayerTypes<CustomLayers...> >& Model()
+  {
+    if (model)
+    {
+      return network;
+    }
+
+    return empty;
+  }
+
+  //! Get the parameters.
+  OutputDataType const& Parameters() const { return weights; }
+  //! Modify the parameters.
+  OutputDataType& Parameters() { return weights; }
+
+  //! Get the input parameter.
+  InputDataType const& InputParameter() const { return inputParameter; }
+  //! Modify the input parameter.
+  InputDataType& InputParameter() { return inputParameter; }
+
+  //! Get the output parameter.
+  OutputDataType const& OutputParameter() const { return outputParameter; }
+  //! Modify the output parameter.
+  OutputDataType& OutputParameter() { return outputParameter; }
+
+  //! Get the delta.
+  OutputDataType const& Delta() const { return delta; }
+  //! Modify the delta.
+  OutputDataType& Delta() { return delta; }
+
+  //! Get the gradient.
+  OutputDataType const& Gradient() const { return gradient; }
+  //! Modify the gradient.
+  OutputDataType& Gradient() { return gradient; }
+
+  /**
+   * Serialize the layer.
+   */
+  template<typename Archive>
+  void serialize(Archive& ar, const unsigned int /* version */);
+
+ private:
+  //! Locally-stored number of input units.
+  size_t inSize;
+
+  //! Parameter which indicates if the modules should be exposed.
+  bool model;
+
+  //! Indicator if we already initialized the model.
+  bool reset;
+
+  //! Locally-stored network modules.
+  std::vector<LayerTypes<CustomLayers...> > network;
+
+  //! Locally-stored empty list of modules.
+  std::vector<LayerTypes<CustomLayers...> > empty;
+
+  //! Locally-stored weight object.
+  OutputDataType weights;
+
+  //! Locally-stored delta object.
+  OutputDataType delta;
+
+  //! Locally-stored gradient object.
+  OutputDataType gradient;
+
+  //! Weights for transformation of output.
+  OutputDataType transformWeight;
+
+  //! Bias for transformation of output.
+  OutputDataType transformBias;
+
+  //! Locally-stored transform gate parameters.
+  OutputDataType transformGate;
+
+  //! Locally-stored transform gate activation.
+  OutputDataType transformGateActivation;
+
+  //! Locally-stored transform gate error.
+  OutputDataType transformGateError;
+
+  //! Locally-stored input parameter object.
+  InputDataType inputParameter;
+
+  //! Locally-stored output parameter object.
+  OutputDataType outputParameter;
+
+  //! The input width.
+  size_t width;
+
+  //! The input height.
+  size_t height;
+
+  //! The normal output without highway network.
+  OutputDataType networkOutput;
+
+  //! Locally-stored delta visitor.
+  DeltaVisitor deltaVisitor;
+
+  //! Locally-stored output parameter visitor.
+  OutputParameterVisitor outputParameterVisitor;
+
+  //! Locally-stored delete visitor.
+  DeleteVisitor deleteVisitor;
+
+  //! Locally-stored output width visitor.
+  OutputWidthVisitor outputWidthVisitor;
+
+  //! Locally-stored output height visitor.
+  OutputHeightVisitor outputHeightVisitor;
+}; // class Highway
+
+} // namespace ann
+} // namespace mlpack
+
+// Include implementation.
+#include "highway_impl.hpp"
+
+#endif

src/mlpack/methods/ann/layer/highway_impl.hpp

@@ -0,0 +1,249 @@
+/**
+ * @file highway_impl.hpp
+ * @author Konstantin Sidorov
+ * @author Saksham Bansal
+ *
+ * Implementation of Highway layer.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_METHODS_ANN_LAYER_HIGHWAY_IMPL_HPP
+#define MLPACK_METHODS_ANN_LAYER_HIGHWAY_IMPL_HPP
+
+// In case it hasn't yet been included.
+#include "highway.hpp"
+
+#include "../visitor/forward_visitor.hpp"
+#include "../visitor/backward_visitor.hpp"
+#include "../visitor/gradient_visitor.hpp"
+#include "../visitor/set_input_height_visitor.hpp"
+#include "../visitor/set_input_width_visitor.hpp"
+
+namespace mlpack {
+namespace ann /** Artificial Neural Network. */ {
+
+template<typename InputDataType, typename OutputDataType,
+         typename... CustomLayers>
+Highway<InputDataType, OutputDataType, CustomLayers...>::Highway() :
+    inSize(0),
+    model(true),
+    reset(false),
+    width(0),
+    height(0)
+{
+  // Nothing to do here.
+}
+
+template<
+    typename InputDataType, typename OutputDataType, typename... CustomLayers>
+Highway<InputDataType, OutputDataType, CustomLayers...>::Highway(
+    const size_t inSize,
+    const bool model) :
+    inSize(inSize),
+    model(model),
+    reset(false),
+    width(0),
+    height(0)
+{
+  weights.set_size(inSize * inSize + inSize, 1);
+}
+
+template<typename InputDataType, typename OutputDataType,
+         typename... CustomLayers>
+Highway<InputDataType, OutputDataType, CustomLayers...>::~Highway()
+{
+  if (!model)
+  {
+    for (LayerTypes<CustomLayers...>& layer : network)
+    {
+      boost::apply_visitor(deleteVisitor, layer);
+    }
+  }
+}
+
+template<typename InputDataType, typename OutputDataType,
+         typename... CustomLayers>
+void Highway<
+    InputDataType, OutputDataType, CustomLayers...>::DeleteModules()
+{
+  if (model)
+  {
+    for (LayerTypes<CustomLayers...>& layer : network)
+    {
+      boost::apply_visitor(deleteVisitor, layer);
+    }
+  }
+}
+
+template<typename InputDataType, typename OutputDataType,
+         typename... CustomLayers>
+void Highway<InputDataType, OutputDataType, CustomLayers...>::Reset()
+{
+  transformWeight = arma::mat(weights.memptr(), inSize, inSize, false, false);
+  transformBias = arma::mat(weights.memptr() + transformWeight.n_elem,
+      inSize, 1, false, false);
+}
+
+template<typename InputDataType, typename OutputDataType,
+         typename... CustomLayers>
+template<typename eT>
+void Highway<InputDataType, OutputDataType, CustomLayers...>::Forward(
+    arma::Mat<eT>&& input, arma::Mat<eT>&& output)
+{
+  boost::apply_visitor(ForwardVisitor(std::move(input), std::move(
+      boost::apply_visitor(outputParameterVisitor, network.front()))),
+      network.front());
+
+  if (!reset)
+  {
+    if (boost::apply_visitor(outputWidthVisitor, network.front()) != 0)
+    {
+      width = boost::apply_visitor(outputWidthVisitor, network.front());
+    }
+
+    if (boost::apply_visitor(outputHeightVisitor, network.front()) != 0)
+    {
+      height = boost::apply_visitor(outputHeightVisitor, network.front());
+    }
+  }
+
+  for (size_t i = 1; i < network.size(); ++i)
+  {
+    if (!reset)
+    {
+      // Set the input width.
+      boost::apply_visitor(SetInputWidthVisitor(width), network[i]);
+
+      // Set the input height.
+      boost::apply_visitor(SetInputHeightVisitor(height), network[i]);
+    }
+
+    boost::apply_visitor(ForwardVisitor(std::move(boost::apply_visitor(
+        outputParameterVisitor, network[i - 1])), std::move(
+        boost::apply_visitor(outputParameterVisitor, network[i]))),
+        network[i]);
+
+    if (!reset)
+    {
+      // Get the output width.
+      if (boost::apply_visitor(outputWidthVisitor, network[i]) != 0)
+      {
+        width = boost::apply_visitor(outputWidthVisitor, network[i]);
+      }
+
+      // Get the output height.
+      if (boost::apply_visitor(outputHeightVisitor, network[i]) != 0)
+      {
+        height = boost::apply_visitor(outputHeightVisitor, network[i]);
+      }
+    }
+  }
+  if (!reset)
+  {
+    reset = true;
+  }
+
+  output = boost::apply_visitor(outputParameterVisitor, network.back());
+
+  if (arma::size(output) != arma::size(input))
+  {
+    Log::Fatal << "The sizes of the output and input matrices of the Highway"
+        << " network should be equal. Please examine the network layers.";
+  }
+
+  transformGate = transformWeight * input;
+  transformGate.each_col() += transformBias;
+  transformGateActivation = 1.0 /(1 + arma::exp(-transformGate));
+  inputParameter = input;
+  networkOutput = output;
+  output = (output % transformGateActivation) +
+      (input % (1 - transformGateActivation));
+}
+
+template<typename InputDataType, typename OutputDataType,
+         typename... CustomLayers>
+template<typename eT>
+void Highway<InputDataType, OutputDataType, CustomLayers...>::Backward(
+    const arma::Mat<eT>&& /* input */,
+    arma::Mat<eT>&& gy,
+    arma::Mat<eT>&& g)
+{
+  boost::apply_visitor(BackwardVisitor(std::move(boost::apply_visitor(
+      outputParameterVisitor, network.back())),
+      std::move(gy % transformGateActivation),
+      std::move(boost::apply_visitor(deltaVisitor, network.back()))),
+      network.back());
+
+  for (size_t i = 2; i < network.size() + 1; ++i)
+  {
+    boost::apply_visitor(BackwardVisitor(std::move(boost::apply_visitor(
+        outputParameterVisitor, network[network.size() - i])), std::move(
+        boost::apply_visitor(deltaVisitor, network[network.size() - i + 1])),
+        std::move(boost::apply_visitor(deltaVisitor,
+        network[network.size() - i]))), network[network.size() - i]);
+  }
+
+  g = boost::apply_visitor(deltaVisitor, network.front());
+
+  transformGateError = gy % (networkOutput - inputParameter) %
+      transformGateActivation % (1.0 - transformGateActivation);
+  g += transformWeight.t() * transformGateError;
+  g += gy % (1 - transformGateActivation);
+}
+
+template<typename InputDataType, typename OutputDataType,
+         typename... CustomLayers>
+template<typename eT>
+void Highway<InputDataType, OutputDataType, CustomLayers...>::Gradient(
+    arma::Mat<eT>&& input,
+    arma::Mat<eT>&& error,
+    arma::Mat<eT>&& gradient)
+{
+  boost::apply_visitor(GradientVisitor(std::move(boost::apply_visitor(
+      outputParameterVisitor, network[network.size() - 2])),
+      std::move(error % transformGateActivation)), network.back());
+
+  for (size_t i = 2; i < network.size(); ++i)
+  {
+    boost::apply_visitor(GradientVisitor(std::move(boost::apply_visitor(
+        outputParameterVisitor, network[network.size() - i - 1])), std::move(
+        boost::apply_visitor(deltaVisitor, network[network.size() - i + 1]))),
+        network[network.size() - i]);
+  }
+
+  boost::apply_visitor(GradientVisitor(std::move(input), std::move(
+      boost::apply_visitor(deltaVisitor, network[1]))), network.front());
+
+  gradient.submat(0, 0, transformWeight.n_elem - 1, 0) = arma::vectorise(
+      transformGateError * input.t());
+  gradient.submat(transformWeight.n_elem, 0, gradient.n_elem - 1, 0) =
+      arma::sum(transformGateError, 1);
+}
+
+template<typename InputDataType, typename OutputDataType,
+         typename... CustomLayers>
+template<typename Archive>
+void Highway<InputDataType, OutputDataType, CustomLayers...>::serialize(
+    Archive& ar, const unsigned int /* version */)
+{
+  // If loading, delete the old layers and set size for weights.
+  if (Archive::is_loading::value)
+  {
+    for (LayerTypes<CustomLayers...>& layer : network)
+    {
+      boost::apply_visitor(deleteVisitor, layer);
+    }
+    weights.set_size(inSize * inSize + inSize, 1);
+  }
+
+  ar & BOOST_SERIALIZATION_NVP(model);
+  ar & BOOST_SERIALIZATION_NVP(network);
+}
+
+} // namespace ann
+} // namespace mlpack
+
+#endif

src/mlpack/methods/ann/layer/layer.hpp

@@ -19,6 +19,7 @@
 #include "convolution.hpp"
 #include "dropconnect.hpp"
 #include "glimpse.hpp"
+#include "highway.hpp"
 #include "layer_norm.hpp"
 #include "layer_types.hpp"
 #include "linear.hpp"

src/mlpack/methods/ann/layer/layer_norm.hpp

@@ -184,6 +184,9 @@ class LayerNorm
 
   //! Locally-stored normalized input.
   OutputDataType normalized;
+
+  //! Locally-stored zero mean input.
+  OutputDataType inputMean;
 }; // class LayerNorm
 
 } // namespace ann

src/mlpack/methods/ann/layer/layer_norm_impl.hpp

@@ -22,6 +22,7 @@ namespace ann { /** Artificial Neural Network. */
 
 template<typename InputDataType, typename OutputDataType>
 LayerNorm<InputDataType, OutputDataType>::LayerNorm() :
+    size(0),
     eps(1e-8),
     loading(false)
 {
@@ -63,7 +64,7 @@ void LayerNorm<InputDataType, OutputDataType>::Forward(
 
   // Normalize the input.
   output = input.each_row() - mean;
-
+  inputMean = output;
   output.each_row() /= arma::sqrt(variance + eps);
 
   // Reused in the backward and gradient step.
@@ -79,7 +80,6 @@ template<typename eT>
 void LayerNorm<InputDataType, OutputDataType>::Backward(
     const arma::Mat<eT>&& input, arma::Mat<eT>&& gy, arma::Mat<eT>&& g)
 {
-  const arma::mat inputMean = input.each_row() - mean;
   const arma::mat stdInv = 1.0 / arma::sqrt(variance + eps);
 
   // dl / dxhat
@@ -96,8 +96,7 @@ void LayerNorm<InputDataType, OutputDataType>::Backward(
 
   // sum (dl / dxhat * -1 / stdInv) + variance *
   // (sum -2 * (x - mu)) / m.
-  g.each_row() += (arma::sum(norm.each_row() % -stdInv, 0) + (var %
-      arma::mean(-2 * inputMean, 0))) / input.n_rows;
+  g.each_row() += arma::sum(norm.each_row() % -stdInv, 0) / input.n_rows;
 }
 
 template<typename InputDataType, typename OutputDataType>

src/mlpack/methods/ann/layer/layer_types.hpp

@@ -82,6 +82,12 @@ template<typename InputDataType,
 >
 class Sequential;
 
+template<typename InputDataType,
+         typename OutputDataType,
+         typename... CustomLayers
+>
+class Highway;
+
 template<typename InputDataType,
          typename OutputDataType,
          typename... CustomLayers
@@ -173,6 +179,7 @@ using LayerTypes = boost::variant<
     FlexibleReLU<arma::mat, arma::mat>*,
     Glimpse<arma::mat, arma::mat>*,
     HardTanH<arma::mat, arma::mat>*,
+    Highway<arma::mat, arma::mat>*,
     Join<arma::mat, arma::mat>*,
     LayerNorm<arma::mat, arma::mat>*,
     LeakyReLU<arma::mat, arma::mat>*,
@@ -191,14 +198,16 @@ using LayerTypes = boost::variant<
     NegativeLogLikelihood<arma::mat, arma::mat>*,
     PReLU<arma::mat, arma::mat>*,
     Recurrent<arma::mat, arma::mat>*,
-    RecurrentAttention<arma::mat, arma::mat>*,
+    // TODO find workaround to support more than 50 types
+    // as boost::variant can only be used for up to 50 types.
+    // RecurrentAttention<arma::mat, arma::mat>*,
     ReinforceNormal<arma::mat, arma::mat>*,
     Reparametrization<arma::mat, arma::mat>*,
     Select<arma::mat, arma::mat>*,
     Sequential<arma::mat, arma::mat, false>*,
     Sequential<arma::mat, arma::mat, true>*,
     Subview<arma::mat, arma::mat>*,
-    VRClassReward<arma::mat, arma::mat>*,
+    // VRClassReward<arma::mat, arma::mat>*,
     CustomLayers*...
 >;
 

src/mlpack/methods/ann/layer/max_pooling_impl.hpp

@@ -148,6 +148,11 @@ void MaxPooling<InputDataType, OutputDataType>::serialize(
   ar & BOOST_SERIALIZATION_NVP(dW);
   ar & BOOST_SERIALIZATION_NVP(dH);
   ar & BOOST_SERIALIZATION_NVP(batchSize);
+  ar & BOOST_SERIALIZATION_NVP(floor);
+  ar & BOOST_SERIALIZATION_NVP(inputWidth);
+  ar & BOOST_SERIALIZATION_NVP(inputHeight);
+  ar & BOOST_SERIALIZATION_NVP(outputWidth);
+  ar & BOOST_SERIALIZATION_NVP(outputHeight);
 }
 
 } // namespace ann

src/mlpack/methods/ann/layer/mean_pooling_impl.hpp

@@ -122,6 +122,11 @@ void MeanPooling<InputDataType, OutputDataType>::serialize(
   ar & BOOST_SERIALIZATION_NVP(dW);
   ar & BOOST_SERIALIZATION_NVP(dH);
   ar & BOOST_SERIALIZATION_NVP(batchSize);
+  ar & BOOST_SERIALIZATION_NVP(floor);
+  ar & BOOST_SERIALIZATION_NVP(inputWidth);
+  ar & BOOST_SERIALIZATION_NVP(inputHeight);
+  ar & BOOST_SERIALIZATION_NVP(outputWidth);
+  ar & BOOST_SERIALIZATION_NVP(outputHeight);
 }
 
 } // namespace ann

src/mlpack/methods/decision_tree/decision_tree.hpp

@@ -64,6 +64,7 @@ class DecisionTree :
    * @param numClasses Number of classes in the dataset.
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    */
   template<typename MatType, typename LabelsType>
@@ -73,6 +74,7 @@ class DecisionTree :
                const size_t numClasses,
                const size_t minimumLeafSize = 10,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -89,6 +91,7 @@ class DecisionTree :
    * @param numClasses Number of classes in the dataset.
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    */
   template<typename MatType, typename LabelsType>
@@ -97,6 +100,7 @@ class DecisionTree :
                const size_t numClasses,
                const size_t minimumLeafSize = 10,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -116,6 +120,7 @@ class DecisionTree :
    * @param weights The weight list of given label.
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    */
   template<typename MatType, typename LabelsType, typename WeightsType>
@@ -126,6 +131,7 @@ class DecisionTree :
                WeightsType weights,
                const size_t minimumLeafSize = 10,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType(),
                const std::enable_if_t<arma::is_arma_type<
@@ -147,6 +153,7 @@ class DecisionTree :
    * @param weights The Weight list of given labels.
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    */
   template<typename MatType, typename LabelsType, typename WeightsType>
@@ -156,6 +163,7 @@ class DecisionTree :
                WeightsType weights,
                const size_t minimumLeafSize = 10,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType(),
                const std::enable_if_t<arma::is_arma_type<
@@ -222,6 +230,7 @@ class DecisionTree :
    * @param weights Weights of all the labels
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @return The final entropy of decision tree.
    */
@@ -232,6 +241,7 @@ class DecisionTree :
                const size_t numClasses,
                const size_t minimumLeafSize = 10,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -249,6 +259,7 @@ class DecisionTree :
    * @param weights Weights of all the labels
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @return The final entropy of decision tree.
    */
@@ -258,6 +269,7 @@ class DecisionTree :
                const size_t numClasses,
                const size_t minimumLeafSize = 10,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -278,6 +290,7 @@ class DecisionTree :
    * @param weights Weights of all the labels
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @return The final entropy of decision tree.
    */
@@ -289,6 +302,7 @@ class DecisionTree :
                WeightsType weights,
                const size_t minimumLeafSize = 10,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType(),
                const std::enable_if_t<arma::is_arma_type<typename
@@ -309,6 +323,7 @@ class DecisionTree :
    * @param weights Weights of all the labels
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @return The final entropy of decision tree.
    */
@@ -319,6 +334,7 @@ class DecisionTree :
                WeightsType weights,
                const size_t minimumLeafSize = 10,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType(),
                const std::enable_if_t<arma::is_arma_type<typename
@@ -453,6 +469,7 @@ class DecisionTree :
    * @param numClasses Number of classes in the dataset.
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @return The final entropy of decision tree.
    */
   template<bool UseWeights, typename MatType>
@@ -465,6 +482,7 @@ class DecisionTree :
                arma::rowvec& weights,
                const size_t minimumLeafSize,
                const double minimumGainSplit,
+               const size_t maximumDepth,
                DimensionSelectionType& dimensionSelector);
 
   /**
@@ -480,6 +498,7 @@ class DecisionTree :
    * @param numClasses Number of classes in the dataset.
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for the node to split.
+   * @param maximumDepth Maximum depth for the tree.
    * @return The final entropy of decision tree.
    */
   template<bool UseWeights, typename MatType>
@@ -491,6 +510,7 @@ class DecisionTree :
                arma::rowvec& weights,
                const size_t minimumLeafSize,
                const double minimumGainSplit,
+               const size_t maximumDepth,
                DimensionSelectionType& dimensionSelector);
 };
 

src/mlpack/methods/decision_tree/decision_tree_impl.hpp

@@ -35,6 +35,7 @@ DecisionTree<FitnessFunction,
     const size_t numClasses,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType dimensionSelector)
 {
   using TrueMatType = typename std::decay<MatType>::type;
@@ -50,7 +51,8 @@ DecisionTree<FitnessFunction,
   // Pass off work to the Train() method.
   arma::rowvec weights; // Fake weights, not used.
   Train<false>(tmpData, 0, tmpData.n_cols, datasetInfo, tmpLabels, numClasses,
-      weights, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      weights, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 //! Construct and train.
@@ -72,6 +74,7 @@ DecisionTree<FitnessFunction,
     const size_t numClasses,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType dimensionSelector)
 {
   using TrueMatType = typename std::decay<MatType>::type;
@@ -87,7 +90,7 @@ DecisionTree<FitnessFunction,
   // Pass off work to the Train() method.
   arma::rowvec weights; // Fake weights, not used.
   Train<false>(tmpData, 0, tmpData.n_cols, tmpLabels, numClasses, weights,
-      minimumLeafSize, minimumGainSplit, dimensionSelector);
+      minimumLeafSize, minimumGainSplit, maximumDepth, dimensionSelector);
 }
 
 //! Construct and train with weights.
@@ -111,6 +114,7 @@ DecisionTree<FitnessFunction,
     WeightsType weights,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType dimensionSelector,
     const std::enable_if_t<
         arma::is_arma_type<
@@ -131,7 +135,8 @@ DecisionTree<FitnessFunction,
 
   // Pass off work to the weighted Train() method.
   Train<true>(tmpData, 0, tmpData.n_cols, datasetInfo, tmpLabels, numClasses,
-      tmpWeights, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      tmpWeights, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 //! Construct and train with weights.
@@ -154,6 +159,7 @@ DecisionTree<FitnessFunction,
     WeightsType weights,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType dimensionSelector,
     const std::enable_if_t<
         arma::is_arma_type<
@@ -174,7 +180,7 @@ DecisionTree<FitnessFunction,
 
   // Pass off work to the weighted Train() method.
   Train<true>(tmpData, 0, tmpData.n_cols, tmpLabels, numClasses, tmpWeights,
-      minimumLeafSize, minimumGainSplit, dimensionSelector);
+      minimumLeafSize, minimumGainSplit, maximumDepth, dimensionSelector);
 }
 
 //! Construct, don't train.
@@ -372,6 +378,7 @@ double DecisionTree<FitnessFunction,
     const size_t numClasses,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType dimensionSelector)
 {
   // Sanity check on data.
@@ -397,7 +404,7 @@ double DecisionTree<FitnessFunction,
   // Pass off work to the Train() method.
   arma::rowvec weights; // Fake weights, not used.
   return Train<false>(tmpData, 0, tmpData.n_cols, datasetInfo, tmpLabels,
-      numClasses, weights, minimumLeafSize, minimumGainSplit,
+      numClasses, weights, minimumLeafSize, minimumGainSplit, maximumDepth,
       dimensionSelector);
 }
 
@@ -420,6 +427,7 @@ double DecisionTree<FitnessFunction,
     const size_t numClasses,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType dimensionSelector)
 {
   // Sanity check on data.
@@ -445,7 +453,8 @@ double DecisionTree<FitnessFunction,
   // Pass off work to the Train() method.
   arma::rowvec weights; // Fake weights, not used.
   return Train<false>(tmpData, 0, tmpData.n_cols, tmpLabels, numClasses,
-      weights, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      weights, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 //! Train on the given weighted data.
@@ -469,6 +478,7 @@ double DecisionTree<FitnessFunction,
     WeightsType weights,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType dimensionSelector,
     const std::enable_if_t<
         arma::is_arma_type<
@@ -499,7 +509,7 @@ double DecisionTree<FitnessFunction,
 
   // Pass off work to the Train() method.
   return Train<true>(tmpData, 0, tmpData.n_cols, datasetInfo, tmpLabels,
-      numClasses, tmpWeights, minimumLeafSize, minimumGainSplit,
+      numClasses, tmpWeights, minimumLeafSize, minimumGainSplit, maximumDepth,
       dimensionSelector);
 }
 
@@ -523,6 +533,7 @@ double DecisionTree<FitnessFunction,
     WeightsType weights,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType dimensionSelector,
     const std::enable_if_t<
         arma::is_arma_type<
@@ -553,7 +564,8 @@ double DecisionTree<FitnessFunction,
 
   // Pass off work to the Train() method.
   return Train<true>(tmpData, 0, tmpData.n_cols, tmpLabels, numClasses,
-      tmpWeights, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      tmpWeights, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 //! Train on the given data.
@@ -579,6 +591,7 @@ double DecisionTree<FitnessFunction,
     arma::rowvec& weights,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType& dimensionSelector)
 {
   // Clear children if needed.
@@ -597,48 +610,52 @@ double DecisionTree<FitnessFunction,
       UseWeights ? weights.subvec(begin, begin + count - 1) : weights);
   size_t bestDim = datasetInfo.Dimensionality(); // This means "no split".
   const size_t end = dimensionSelector.End();
-  for (size_t i = dimensionSelector.Begin(); i != end;
-       i = dimensionSelector.Next())
+
+  if (maximumDepth != 1)
   {
-    double dimGain = -DBL_MAX;
-    if (datasetInfo.Type(i) == data::Datatype::categorical)
+    for (size_t i = dimensionSelector.Begin(); i != end;
+         i = dimensionSelector.Next())
     {
-      dimGain = CategoricalSplit::template SplitIfBetter<UseWeights>(bestGain,
-          data.cols(begin, begin + count - 1).row(i),
-          datasetInfo.NumMappings(i),
-          labels.subvec(begin, begin + count - 1),
-          numClasses,
-          UseWeights ? weights.subvec(begin, begin + count - 1) : weights,
-          minimumLeafSize,
-          minimumGainSplit,
-          classProbabilities,
-          *this);
+      double dimGain = -DBL_MAX;
+      if (datasetInfo.Type(i) == data::Datatype::categorical)
+      {
+        dimGain = CategoricalSplit::template SplitIfBetter<UseWeights>(bestGain,
+            data.cols(begin, begin + count - 1).row(i),
+            datasetInfo.NumMappings(i),
+            labels.subvec(begin, begin + count - 1),
+            numClasses,
+            UseWeights ? weights.subvec(begin, begin + count - 1) : weights,
+            minimumLeafSize,
+            minimumGainSplit,
+            classProbabilities,
+            *this);
+      }
+      else if (datasetInfo.Type(i) == data::Datatype::numeric)
+      {
+        dimGain = NumericSplit::template SplitIfBetter<UseWeights>(bestGain,
+            data.cols(begin, begin + count - 1).row(i),
+            labels.subvec(begin, begin + count - 1),
+            numClasses,
+            UseWeights ? weights.subvec(begin, begin + count - 1) : weights,
+            minimumLeafSize,
+            minimumGainSplit,
+            classProbabilities,
+            *this);
+      }
+
+      // If the splitter reported that it did not split, move to the next
+      // dimension.
+      if (dimGain == DBL_MAX)
+        continue;
+
+      // Was there an improvement?  If so mark that it's the new best dimension.
+      bestDim = i;
+      bestGain = dimGain;
+
+      // If the gain is the best possible, no need to keep looking.
+      if (bestGain >= 0.0)
+        break;
     }
-    else if (datasetInfo.Type(i) == data::Datatype::numeric)
-    {
-      dimGain = NumericSplit::template SplitIfBetter<UseWeights>(bestGain,
-          data.cols(begin, begin + count - 1).row(i),
-          labels.subvec(begin, begin + count - 1),
-          numClasses,
-          UseWeights ? weights.subvec(begin, begin + count - 1) : weights,
-          minimumLeafSize,
-          minimumGainSplit,
-          classProbabilities,
-          *this);
-    }
-
-    // If the splitter reported that it did not split, move to the next
-    // dimension.
-    if (dimGain == DBL_MAX)
-      continue;
-
-    // Was there an improvement?  If so mark that it's the new best dimension.
-    bestDim = i;
-    bestGain = dimGain;
-
-    // If the gain is the best possible, no need to keep looking.
-    if (bestGain >= 0.0)
-      break;
   }
 
   // Did we split or not?  If so, then split the data and create the children.
@@ -707,14 +724,15 @@ double DecisionTree<FitnessFunction,
         child->Train<UseWeights>(data, currentChildBegin,
             currentCol - currentChildBegin, datasetInfo, labels, numClasses,
             weights, currentCol - currentChildBegin, minimumGainSplit,
-            dimensionSelector);
+            maximumDepth - 1, dimensionSelector);
       }
       else
       {
         // During recursion entropy of child node may change.
         double childGain = child->Train<UseWeights>(data, currentChildBegin,
             currentCol - currentChildBegin, datasetInfo, labels, numClasses,
-            weights, minimumLeafSize, minimumGainSplit, dimensionSelector);
+            weights, minimumLeafSize, minimumGainSplit, maximumDepth - 1,
+            dimensionSelector);
         bestGain += double(childCounts[i]) / double(count) * (-childGain);
       }
       children.push_back(child);
@@ -732,6 +750,7 @@ double DecisionTree<FitnessFunction,
         numClasses,
         UseWeights ? weights.subvec(begin, begin + count - 1) : weights);
   }
+
   return -bestGain;
 }
 
@@ -757,6 +776,7 @@ double DecisionTree<FitnessFunction,
     arma::rowvec& weights,
     const size_t minimumLeafSize,
     const double minimumGainSplit,
+    const size_t maximumDepth,
     DimensionSelectionType& dimensionSelector)
 {
   // Clear children if needed.
@@ -777,33 +797,37 @@ double DecisionTree<FitnessFunction,
       numClasses,
       UseWeights ? weights.subvec(begin, begin + count - 1) : weights);
   size_t bestDim = data.n_rows; // This means "no split".
-  for (size_t i = dimensionSelector.Begin(); i != dimensionSelector.End();
-       i = dimensionSelector.Next())
+
+  if (maximumDepth != 1)
   {
-    const double dimGain = NumericSplitType<FitnessFunction>::template
-        SplitIfBetter<UseWeights>(bestGain,
-                                  data.cols(begin, begin + count - 1).row(i),
-                                  labels.cols(begin, begin + count - 1),
-                                  numClasses,
-                                  UseWeights ?
-                                      weights.cols(begin, begin + count - 1) :
-                                      weights,
-                                  minimumLeafSize,
-                                  minimumGainSplit,
-                                  classProbabilities,
-                                  *this);
+    for (size_t i = dimensionSelector.Begin(); i != dimensionSelector.End();
+         i = dimensionSelector.Next())
+    {
+      const double dimGain = NumericSplitType<FitnessFunction>::template
+          SplitIfBetter<UseWeights>(bestGain,
+                                    data.cols(begin, begin + count - 1).row(i),
+                                    labels.cols(begin, begin + count - 1),
+                                    numClasses,
+                                    UseWeights ?
+                                        weights.cols(begin, begin + count - 1) :
+                                        weights,
+                                    minimumLeafSize,
+                                    minimumGainSplit,
+                                    classProbabilities,
+                                    *this);
 
-    // If the splitter did not report that it improved, then move to the next
-    // dimension.
-    if (dimGain == DBL_MAX)
-      continue;
+      // If the splitter did not report that it improved, then move to the next
+      // dimension.
+      if (dimGain == DBL_MAX)
+        continue;
 
-    bestDim = i;
-    bestGain = dimGain;
+      bestDim = i;
+      bestGain = dimGain;
 
-    // If the gain is the best possible, no need to keep looking.
-    if (bestGain >= 0.0)
-      break;
+      // If the gain is the best possible, no need to keep looking.
+      if (bestGain >= 0.0)
+        break;
+    }
   }
 
   // Did we split or not?  If so, then split the data and create the children.
@@ -858,7 +882,7 @@ double DecisionTree<FitnessFunction,
       {
         child->Train<UseWeights>(data, currentChildBegin,
             currentCol - currentChildBegin, labels, numClasses, weights,
-            currentCol - currentChildBegin, minimumGainSplit,
+            currentCol - currentChildBegin, minimumGainSplit, maximumDepth - 1,
             dimensionSelector);
       }
       else
@@ -866,7 +890,8 @@ double DecisionTree<FitnessFunction,
         // During recursion entropy of child node may change.
         double childGain = child->Train<UseWeights>(data, currentChildBegin,
             currentCol - currentChildBegin, labels, numClasses, weights,
-            minimumLeafSize, minimumGainSplit, dimensionSelector);
+            minimumLeafSize, minimumGainSplit, maximumDepth - 1,
+            dimensionSelector);
         bestGain += double(childCounts[i]) / double(count) * (-childGain);
       }
       children.push_back(child);
@@ -883,6 +908,7 @@ double DecisionTree<FitnessFunction,
         numClasses,
         UseWeights ? weights.subvec(begin, begin + count - 1) : weights);
   }
+
   return -bestGain;
 }
 

src/mlpack/methods/decision_tree/decision_tree_main.cpp

@@ -47,7 +47,9 @@ PROGRAM_INFO("Decision tree",
     " parameter specifies the minimum number of training points that must fall"
     " into each leaf for it to be split.  The " +
     PRINT_PARAM_STRING("minimum_gain_split") + " parameter specifies "
-    "the minimum gain that is needed for the node to split. If " +
+    "the minimum gain that is needed for the node to split.  The " +
+    PRINT_PARAM_STRING("maximum_depth") + " parameter specifies "
+    "the maximum depth of the tree.  If " +
     PRINT_PARAM_STRING("print_training_error") + " is specified, the training "
     "error will be printed."
     "\n\n"
@@ -100,6 +102,8 @@ PARAM_INT_IN("minimum_leaf_size", "Minimum number of points in a leaf.", "n",
     20);
 PARAM_DOUBLE_IN("minimum_gain_split", "Minimum gain for node splitting.", "g",
     1e-7);
+PARAM_INT_IN("maximum_depth", "Maximum depth of the tree (0 means no limit).",
+    "D", 0);
 // This is deprecated and should be removed in mlpack 4.0.0.
 PARAM_FLAG("print_training_error", "Print the training error (deprecated; will "
       "be removed in mlpack 4.0.0).", "e");
@@ -157,6 +161,9 @@ static void mlpackMain()
   RequireParamValue<int>("minimum_leaf_size", [](int x) { return x > 0; }, true,
       "leaf size must be positive");
 
+  RequireParamValue<int>("maximum_depth", [](int x) { return x >= 0; }, true,
+      "maximum depth must not be negative");
+
   RequireParamValue<double>("minimum_gain_split", [](double x)
                          { return (x > 0.0 && x < 1.0); }, true,
                          "gain split must be a fraction in range [0,1]");
@@ -195,6 +202,7 @@ static void mlpackMain()
 
     // Now build the tree.
     const size_t minLeafSize = (size_t) CLI::GetParam<int>("minimum_leaf_size");
+    const size_t maxDepth = (size_t) CLI::GetParam<int>("maximum_depth");
     const double minimumGainSplit =
                            (double) CLI::GetParam<double>("minimum_gain_split");
 
@@ -207,13 +215,14 @@ static void mlpackMain()
           CLI::HasParam("print_training_accuracy"))
       {
         model->tree = DecisionTree<>(trainingSet, model->info, labels,
-            numClasses, std::move(weights), minLeafSize, minimumGainSplit);
+            numClasses, std::move(weights), minLeafSize, minimumGainSplit,
+            maxDepth);
       }
       else
       {
         model->tree = DecisionTree<>(std::move(trainingSet), model->info,
             std::move(labels), numClasses, std::move(weights), minLeafSize,
-            minimumGainSplit);
+            minimumGainSplit, maxDepth);
       }
     }
     else
@@ -221,12 +230,13 @@ static void mlpackMain()
       if (CLI::HasParam("print_training_error"))
       {
         model->tree = DecisionTree<>(trainingSet, model->info, labels,
-            numClasses, minLeafSize, minimumGainSplit);
+            numClasses, minLeafSize, minimumGainSplit, maxDepth);
       }
       else
       {
         model->tree = DecisionTree<>(std::move(trainingSet), model->info,
-            std::move(labels), numClasses, minLeafSize, minimumGainSplit);
+            std::move(labels), numClasses, minLeafSize, minimumGainSplit,
+            maxDepth);
       }
     }
 

src/mlpack/methods/gmm/gmm_train_main.cpp

@@ -145,12 +145,18 @@ static void mlpackMain()
       "number of Gaussians must be positive");
   const int gaussians = CLI::GetParam<int>("gaussians");
 
+  RequireParamValue<int>("trials", [](int x) { return x > 0; }, true,
+      "trials must be greater than 0");
+
   ReportIgnoredParam({{ "diagonal_covariance", true }}, "no_force_positive");
   RequireAtLeastOnePassed({ "output_model" }, false, "no model will be saved");
 
   RequireParamValue<double>("noise", [](double x) { return x >= 0.0; }, true,
       "variance of noise must be greater than or equal to 0");
 
+  RequireParamValue<int>("max_iterations", [](int x) { return x >= 0; }, true,
+      "max_iterations must be greater than or equal to 0");
+
   arma::mat dataPoints = std::move(CLI::GetParam<arma::mat>("input"));
 
   // Do we need to add noise to the dataset?

src/mlpack/methods/random_forest/random_forest.hpp

@@ -50,6 +50,7 @@ class RandomForest
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each tree's leaf nodes.
    * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    */
   template<typename MatType>
@@ -59,6 +60,7 @@ class RandomForest
                const size_t numTrees = 20,
                const size_t minimumLeafSize = 1,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -77,6 +79,7 @@ class RandomForest
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each tree's leaf nodes.
    * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    */
   template<typename MatType>
@@ -87,6 +90,7 @@ class RandomForest
                const size_t numTrees = 20,
                const size_t minimumLeafSize = 1,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -101,6 +105,9 @@ class RandomForest
    * @param weights Weights (importances) of each point in the dataset.
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each tree's leaf nodes.
+   * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
+   * @param dimensionSelector Instantiated dimension selection policy.
    */
   template<typename MatType>
   RandomForest(const MatType& dataset,
@@ -110,6 +117,7 @@ class RandomForest
                const size_t numTrees = 20,
                const size_t minimumLeafSize = 1,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -129,6 +137,7 @@ class RandomForest
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each tree's leaf nodes.
    * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    */
   template<typename MatType>
@@ -140,6 +149,7 @@ class RandomForest
                const size_t numTrees = 20,
                const size_t minimumLeafSize = 1,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -156,6 +166,7 @@ class RandomForest
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each tree's leaf nodes.
    * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @return The average entropy of all the decision trees trained under forest.
    */
@@ -166,6 +177,7 @@ class RandomForest
                const size_t numTrees = 20,
                const size_t minimumLeafSize = 1,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -185,6 +197,7 @@ class RandomForest
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each tree's leaf nodes.
    * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @return The average entropy of all the decision trees trained under forest.
    */
@@ -196,6 +209,7 @@ class RandomForest
                const size_t numTrees = 20,
                const size_t minimumLeafSize = 1,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -213,6 +227,7 @@ class RandomForest
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each tree's leaf nodes.
    * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @return The average entropy of all the decision trees trained under forest.
    */
@@ -224,6 +239,7 @@ class RandomForest
                const size_t numTrees = 20,
                const size_t minimumLeafSize = 1,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -243,6 +259,7 @@ class RandomForest
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each tree's leaf nodes.
    * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @return The average entropy of all the decision trees trained under forest.
    */
@@ -255,6 +272,7 @@ class RandomForest
                const size_t numTrees = 20,
                const size_t minimumLeafSize = 1,
                const double minimumGainSplit = 1e-7,
+               const size_t maximumDepth = 0,
                DimensionSelectionType dimensionSelector =
                    DimensionSelectionType());
 
@@ -334,6 +352,7 @@ class RandomForest
    * @param numTrees Number of trees in the forest.
    * @param minimumLeafSize Minimum number of points in each leaf node.
    * @param minimumGainSplit Minimum gain for splitting a decision tree node.
+   * @param maximumDepth Maximum depth for the tree.
    * @param dimensionSelector Instantiated dimension selection policy.
    * @tparam UseWeights Whether or not to use the weights parameter.
    * @tparam UseDatasetInfo Whether or not to use the datasetInfo parameter.
@@ -349,6 +368,7 @@ class RandomForest
                const size_t numTrees,
                const size_t minimumLeafSize,
                const double minimumGainSplit,
+               const size_t maximumDepth,
                DimensionSelectionType& dimensionSelector);
 
   //! The trees in the forest.

src/mlpack/methods/random_forest/random_forest_impl.hpp

@@ -38,13 +38,14 @@ RandomForest<
                 const size_t numTrees,
                 const size_t minimumLeafSize,
                 const double minimumGainSplit,
+                const size_t maximumDepth,
                 DimensionSelectionType dimensionSelector)
 {
   // Pass off work to the Train() method.
   data::DatasetInfo info; // Ignored.
   arma::rowvec weights; // Fake weights, not used.
   Train<false, false>(dataset, info, labels, numClasses, weights, numTrees,
-      minimumLeafSize, minimumGainSplit, dimensionSelector);
+      minimumLeafSize, minimumGainSplit, maximumDepth, dimensionSelector);
 }
 
 template<
@@ -68,12 +69,14 @@ RandomForest<
                 const size_t numTrees,
                 const size_t minimumLeafSize,
                 const double minimumGainSplit,
+                const size_t maximumDepth,
                 DimensionSelectionType dimensionSelector)
 {
   // Pass off work to the Train() method.
   arma::rowvec weights; // Fake weights, not used.
   Train<false, true>(dataset, datasetInfo, labels, numClasses, weights,
-      numTrees, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      numTrees, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 template<
@@ -97,12 +100,13 @@ RandomForest<
                 const size_t numTrees,
                 const size_t minimumLeafSize,
                 const double minimumGainSplit,
+                const size_t maximumDepth,
                 DimensionSelectionType dimensionSelector)
 {
   // Pass off work to the Train() method.
   data::DatasetInfo info; // Ignored by Train().
   Train<true, false>(dataset, info, labels, numClasses, weights, numTrees,
-      minimumLeafSize, minimumGainSplit, dimensionSelector);
+      minimumLeafSize, minimumGainSplit, maximumDepth, dimensionSelector);
 }
 
 template<
@@ -127,11 +131,12 @@ RandomForest<
                 const size_t numTrees,
                 const size_t minimumLeafSize,
                 const double minimumGainSplit,
+                const size_t maximumDepth,
                 DimensionSelectionType dimensionSelector)
 {
   // Pass off work to the Train() method.
   Train<true, true>(dataset, datasetInfo, labels, numClasses, weights, numTrees,
-      minimumLeafSize, minimumGainSplit, dimensionSelector);
+      minimumLeafSize, minimumGainSplit, maximumDepth, dimensionSelector);
 }
 
 template<
@@ -154,13 +159,15 @@ double RandomForest<
          const size_t numTrees,
          const size_t minimumLeafSize,
          const double minimumGainSplit,
+         const size_t maximumDepth,
          DimensionSelectionType dimensionSelector)
 {
   // Pass off to Train().
   data::DatasetInfo info; // Ignored by Train().
   arma::rowvec weights; // Ignored by Train().
   return Train<false, false>(dataset, info, labels, numClasses, weights,
-      numTrees, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      numTrees, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 template<
@@ -184,12 +191,14 @@ double RandomForest<
          const size_t numTrees,
          const size_t minimumLeafSize,
          const double minimumGainSplit,
+         const size_t maximumDepth,
          DimensionSelectionType dimensionSelector)
 {
   // Pass off to Train().
   arma::rowvec weights; // Ignored by Train().
   return Train<false, true>(dataset, datasetInfo, labels, numClasses, weights,
-      numTrees, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      numTrees, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 template<
@@ -213,12 +222,14 @@ double RandomForest<
          const size_t numTrees,
          const size_t minimumLeafSize,
          const double minimumGainSplit,
+         const size_t maximumDepth,
          DimensionSelectionType dimensionSelector)
 {
   // Pass off to Train().
   data::DatasetInfo info; // Ignored by Train().
   return Train<false, false>(dataset, info, labels, numClasses, weights,
-      numTrees, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      numTrees, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 template<
@@ -243,11 +254,13 @@ double RandomForest<
          const size_t numTrees,
          const size_t minimumLeafSize,
          const double minimumGainSplit,
+         const size_t maximumDepth,
          DimensionSelectionType dimensionSelector)
 {
   // Pass off to Train().
   return Train<true, true>(dataset, datasetInfo, labels, numClasses, weights,
-      numTrees, minimumLeafSize, minimumGainSplit, dimensionSelector);
+      numTrees, minimumLeafSize, minimumGainSplit, maximumDepth,
+      dimensionSelector);
 }
 
 template<
@@ -447,6 +460,7 @@ double RandomForest<
          const size_t numTrees,
          const size_t minimumLeafSize,
          const double minimumGainSplit,
+         const size_t maximumDepth,
          DimensionSelectionType& dimensionSelector)
 {
   // Train each tree individually.
@@ -472,12 +486,12 @@ double RandomForest<
       {
         avgGain += trees[i].Train(bootstrapDataset, datasetInfo,
             bootstrapLabels, numClasses, bootstrapWeights, minimumLeafSize,
-            minimumGainSplit, dimensionSelector);
+            minimumGainSplit, maximumDepth, dimensionSelector);
       }
       else
       {
         avgGain += trees[i].Train(bootstrapDataset, bootstrapLabels, numClasses,
-            bootstrapWeights, minimumLeafSize, minimumGainSplit,
+            bootstrapWeights, minimumLeafSize, minimumGainSplit, maximumDepth,
             dimensionSelector);
       }
     }
@@ -487,12 +501,12 @@ double RandomForest<
       {
         avgGain += trees[i].Train(bootstrapDataset, datasetInfo,
             bootstrapLabels, numClasses, minimumLeafSize, minimumGainSplit,
-            dimensionSelector);
+            maximumDepth, dimensionSelector);
       }
       else
       {
         avgGain += trees[i].Train(bootstrapDataset, bootstrapLabels, numClasses,
-            minimumLeafSize, minimumGainSplit, dimensionSelector);
+            minimumLeafSize, minimumGainSplit, maximumDepth, dimensionSelector);
       }
     }
     Timer::Stop("train_tree");

src/mlpack/methods/random_forest/random_forest_main.cpp

@@ -51,6 +51,8 @@ PROGRAM_INFO("Random forests",
     PRINT_PARAM_STRING("minimum_gain_split") + " parameter controls the minimum"
     " required gain for a decision tree node to split.  Larger values will "
     "force higher-confidence splits.  The " +
+    PRINT_PARAM_STRING("maximum_depth") + " parameter specifies "
+    "the maximum depth of the tree.  The " +
     PRINT_PARAM_STRING("subspace_dim") + " parameter is used to control the "
     "number of random dimensions chosen for an individual node's split.  If " +
     PRINT_PARAM_STRING("print_training_accuracy") + " is specified, the "
@@ -105,7 +107,8 @@ PARAM_FLAG("print_training_accuracy", "If set, then the accuracy of the model "
 PARAM_INT_IN("num_trees", "Number of trees in the random forest.", "N", 10);
 PARAM_INT_IN("minimum_leaf_size", "Minimum number of points in each leaf "
     "node.", "n", 1);
-
+PARAM_INT_IN("maximum_depth", "Maximum depth of the tree (0 means no limit).",
+    "D", 0);
 PARAM_MATRIX_OUT("probabilities", "Predicted class probabilities for each "
     "point in the test set.", "P");
 PARAM_UROW_OUT("predictions", "Predicted classes for each point in the test "
@@ -177,6 +180,8 @@ static void mlpackMain()
 
   RequireParamValue<int>("minimum_leaf_size", [](int x) { return x > 0; }, true,
       "minimum leaf size must be greater than 0");
+  RequireParamValue<int>("maximum_depth", [](int x) { return x >= 0; }, true,
+      "maximum depth must not be negative");
   RequireParamValue<int>("subspace_dim", [](int x) { return x >= 0; }, true,
       "subspace dimensionality must be nonnegative");
   RequireParamValue<double>("minimum_gain_split",
@@ -205,6 +210,7 @@ static void mlpackMain()
     const size_t numTrees = (size_t) CLI::GetParam<int>("num_trees");
     const size_t minimumLeafSize =
         (size_t) CLI::GetParam<int>("minimum_leaf_size");
+    const size_t maxDepth = (size_t) CLI::GetParam<int>("maximum_depth");
     const double minimumGainSplit = CLI::GetParam<double>("minimum_gain_split");
     const size_t randomDims = (CLI::GetParam<int>("subspace_dim") == 0) ?
         (size_t) std::sqrt(data.n_rows) :
@@ -218,7 +224,7 @@ static void mlpackMain()
 
     // Train the model.
     rfModel->rf.Train(data, labels, numClasses, numTrees, minimumLeafSize,
-        minimumGainSplit, mrds);
+        minimumGainSplit, maxDepth, mrds);
     Timer::Stop("rf_training");
 
     // Did we want training accuracy?

src/mlpack/methods/reinforcement_learning/environment/CMakeLists.txt

@@ -4,6 +4,8 @@ set(SOURCES
   mountain_car.hpp
   cart_pole.hpp
   continuous_mountain_car.hpp
+  multiple_pole_cart.hpp
+  continuous_multiple_pole_cart.hpp
   acrobot.hpp
   pendulum.hpp
   reward_clipping.hpp

src/mlpack/methods/reinforcement_learning/environment/acrobot.hpp

@@ -28,10 +28,10 @@ namespace rl{
 class Acrobot
 {
  public:
-   /*
-    * Implementation of Acrobot State. Each State is a tuple vector
-    * (theta1, thetha2, angular velocity 1, angular velocity 2).
-    */
+  /*
+   * Implementation of Acrobot State. Each State is a tuple vector
+   * (theta1, thetha2, angular velocity 1, angular velocity 2).
+   */
   class State
   {
    public:
@@ -95,21 +95,24 @@ class Acrobot
     size
   };
 
-   /**
-    * Construct a Acrobot instance using the given constants.
-    *
-    * @param gravity The gravity parameter.
-    * @param linkLength1 The length of link 1.
-    * @param linkLength2 The length of link 2.
-    * @param linkMass1 The mass of link 1.
-    * @param linkMass2 The mass of link 2.
-    * @param linkCom1 The position of the center of mass of link 1.
-    * @param linkCom2 The position of the center of mass of link 2.
-    * @param linkMoi The moments of inertia for both link.
-    * @param maxVel1 The max angular velocity of link1.
-    * @param maxVel2 The max angular velocity of link2.
-    * @param dt The differential value.
-    */
+  /**
+   * Construct a Acrobot instance using the given constants.
+   *
+   * @param gravity The gravity parameter.
+   * @param linkLength1 The length of link 1.
+   * @param linkLength2 The length of link 2.
+   * @param linkMass1 The mass of link 1.
+   * @param linkMass2 The mass of link 2.
+   * @param linkCom1 The position of the center of mass of link 1.
+   * @param linkCom2 The position of the center of mass of link 2.
+   * @param linkMoi The moments of inertia for both links.
+   * @param maxVel1 The max angular velocity of link1.
+   * @param maxVel2 The max angular velocity of link2.
+   * @param dt The differential value.
+   * @param doneReward The reward recieved by the agent on success.
+   * @param maxSteps The number of steps after which the episode
+   *    terminates. If the value is 0, there is no limit.
+   */
   Acrobot(const double gravity = 9.81,
           const double linkLength1 = 1.0,
           const double linkLength2 = 1.0,
@@ -121,7 +124,8 @@ class Acrobot
           const double maxVel1 = 4 * M_PI,
           const double maxVel2 = 9 * M_PI,
           const double dt = 0.2,
-          const double doneReward = 0) :
+          const double doneReward = 0,
+          const size_t maxSteps = 0) :
       gravity(gravity),
       linkLength1(linkLength1),
       linkLength2(linkLength2),
@@ -133,7 +137,9 @@ class Acrobot
       maxVel1(maxVel1),
       maxVel2(maxVel2),
       dt(dt),
-      doneReward(doneReward)
+      doneReward(doneReward),
+      maxSteps(maxSteps),
+      stepsPerformed(0)
   { /* Nothing to do here */ }
 
   /**
@@ -147,8 +153,11 @@ class Acrobot
    */
   double Sample(const State& state,
                 const Action& action,
-                State& nextState) const
+                State& nextState)
   {
+    // Update the number of steps performed.
+    stepsPerformed++;
+
     // Make a vector to estimate nextstate.
     arma::colvec currentState = {state.Theta1(), state.Theta2(),
         state.AngularVelocity1(), state.AngularVelocity2()};
@@ -158,19 +167,22 @@ class Acrobot
     nextState.Theta1() = Wrap(currentNextState[0], -M_PI, M_PI);
 
     nextState.Theta2() = Wrap(currentNextState[1], -M_PI, M_PI);
-    //! The value of angular velocity is bounded in min and max value.
 
+    //! The value of angular velocity is bounded in min and max value.
     nextState.AngularVelocity1() = std::min(
         std::max(currentNextState[2], -maxVel1), maxVel1);
     nextState.AngularVelocity2() = std::min(
         std::max(currentNextState[3], -maxVel2), maxVel2);
-    /**
-     * If the acrobot reaches a terminal state, it should be given a positive
-     * reward. This will ensure that the agent learns the goal of the game.
-     */
+
+    // Check if the episode has terminated.
     bool done = IsTerminal(nextState);
-    if (done)
+
+    // Do not reward the agent if time ran out.
+    if (done && maxSteps != 0 && stepsPerformed >= maxSteps)
+      return 0;
+    else if (done)
       return doneReward;
+
     return -1;
   };
 
@@ -183,7 +195,7 @@ class Acrobot
    * @param action The action taken.
    * @param nextState The next state.
    */
-  double Sample(const State& state, const Action& action) const
+  double Sample(const State& state, const Action& action)
   {
     State nextState;
     return Sample(state, action, nextState);
@@ -192,8 +204,9 @@ class Acrobot
   /**
    * This function does random initialization of state space.
    */
-  State InitialSample() const
+  State InitialSample()
   {
+    stepsPerformed = 0;
     return State((arma::randu<arma::colvec>(4) - 0.5) / 5.0);
   }
 
@@ -201,11 +214,23 @@ class Acrobot
    * This function checks if the acrobot has reached the terminal state.
    *
    * @param state The current State.
+   * @return true if state is a terminal state, otherwise false.
    */
   bool IsTerminal(const State& state) const
   {
-    return bool (-std::cos(state.Theta1())-std::cos(state.Theta1() +
-        state.Theta2()) > 1.0);
+    if (maxSteps != 0 && stepsPerformed >= maxSteps)
+    {
+      Log::Info << "Episode terminated due to the maximum number of steps"
+          "being taken.";
+      return true;
+    }
+    else if (-std::cos(state.Theta1()) - std::cos(state.Theta1() +
+        state.Theta2()) > 1.0)
+    {
+      Log::Info << "Episode terminated due to agent succeeding.";
+      return true;
+    }
+    return false;
   }
 
   /**
@@ -295,7 +320,6 @@ class Acrobot
   }
 
   /**
-   *
    * This function calls the RK4 iterative method to estimate the next state
    * based on given ordinary differential equation.
    *
@@ -313,6 +337,14 @@ class Acrobot
     return nextState;
   };
 
+  //! Get the number of steps performed.
+  size_t StepsPerformed() const { return stepsPerformed; }
+
+  //! Get the maximum number of steps allowed.
+  size_t MaxSteps() const { return maxSteps; }
+  //! Set the maximum number of steps allowed.
+  size_t& MaxSteps() { return maxSteps; }
+
  private:
   //! Locally-stored gravity.
   double gravity;
@@ -349,12 +381,13 @@ class Acrobot
 
   //! Locally-stored done reward.
   double doneReward;
-}; // class Acrobot
 
-/**
- * Add an alias for backward compatibility.
- */
-typedef Acrobot Acrobat;
+  //! Locally-stored maximum number of steps.
+  size_t maxSteps;
+
+  //! Locally-stored number of steps performed.
+  size_t stepsPerformed;
+};
 
 } // namespace rl
 } // namespace mlpack

src/mlpack/methods/reinforcement_learning/environment/cart_pole.hpp

@@ -104,6 +104,9 @@ class CartPole
    * @param tau The time interval.
    * @param thetaThresholdRadians The maximum angle.
    * @param xThreshold The maximum position.
+   * @param doneReward Reward recieved by agent on success.
+   * @param maxSteps The number of steps after which the episode
+   *    terminates. If the value is 0, there is no limit.
    */
   CartPole(const double gravity = 9.8,
            const double massCart = 1.0,
@@ -113,7 +116,8 @@ class CartPole
            const double tau = 0.02,
            const double thetaThresholdRadians = 12 * 2 * 3.1416 / 360,
            const double xThreshold = 2.4,
-           const double doneReward = 0.0) :
+           const double doneReward = 0.0,
+           const size_t maxSteps = 0) :
       gravity(gravity),
       massCart(massCart),
       massPole(massPole),
@@ -124,7 +128,9 @@ class CartPole
       tau(tau),
       thetaThresholdRadians(thetaThresholdRadians),
       xThreshold(xThreshold),
-      doneReward(doneReward)
+      doneReward(doneReward),
+      maxSteps(maxSteps),
+      stepsPerformed(0)
   { /* Nothing to do here */ }
 
   /**
@@ -138,8 +144,11 @@ class CartPole
    */
   double Sample(const State& state,
                 const Action& action,
-                State& nextState) const
+                State& nextState)
   {
+    // Update the number of steps performed.
+    stepsPerformed++;
+
     // Calculate acceleration.
     double force = action ? forceMag : -forceMag;
     double cosTheta = std::cos(state.Angle());
@@ -156,13 +165,15 @@ class CartPole
     nextState.Angle() = state.Angle() + tau * state.AngularVelocity();
     nextState.AngularVelocity() = state.AngularVelocity() + tau * thetaAcc;
 
-    /**
-     * It is important to note that if the cartpole is falling down, it should
-     * be penalized.
-     */
+    // Check if the episode has terminated.
     bool done = IsTerminal(nextState);
-    if (done)
+
+    // Do not reward agent if it failed.
+    if (done && maxSteps != 0 && stepsPerformed >= maxSteps)
       return doneReward;
+    else if (done)
+      return 0;
+
     /**
      * When done is false, it means that the cartpole has fallen down.
      * For this case the reward is 1.0.
@@ -178,7 +189,7 @@ class CartPole
    * @param action The current action.
    * @return reward, it's always 1.0.
    */
-  double Sample(const State& state, const Action& action) const
+  double Sample(const State& state, const Action& action)
   {
     State nextState;
     return Sample(state, action, nextState);
@@ -189,23 +200,43 @@ class CartPole
    *
    * @return Initial state for each episode.
    */
-  State InitialSample() const
+  State InitialSample()
   {
+    stepsPerformed = 0;
     return State((arma::randu<arma::colvec>(4) - 0.5) / 10.0);
   }
 
   /**
-   * Whether given state is a terminal state.
+   * This function checks if the cart has reached the terminal state.
    *
    * @param state The desired state.
    * @return true if state is a terminal state, otherwise false.
    */
   bool IsTerminal(const State& state) const
   {
-    return std::abs(state.Position()) > xThreshold ||
-        std::abs(state.Angle()) > thetaThresholdRadians;
+    if (maxSteps != 0 && stepsPerformed >= maxSteps)
+    {
+      Log::Info << "Episode terminated due to the maximum number of steps"
+          "being taken.";
+      return true;
+    }
+    else if (std::abs(state.Position()) > xThreshold ||
+        std::abs(state.Angle()) > thetaThresholdRadians)
+    {
+      Log::Info << "Episode terminated due to agent failing.";
+      return true;
+    }
+    return false;
   }
 
+  //! Get the number of steps performed.
+  size_t StepsPerformed() const { return stepsPerformed; }
+
+  //! Get the maximum number of steps allowed.
+  size_t MaxSteps() const { return maxSteps; }
+  //! Set the maximum number of steps allowed.
+  size_t& MaxSteps() { return maxSteps; }
+
  private:
   //! Locally-stored gravity.
   double gravity;
@@ -239,6 +270,12 @@ class CartPole
 
   //! Locally-stored done reward.
   double doneReward;
+
+  //! Locally-stored maximum number of steps.
+  size_t maxSteps;
+
+  //! Locally-stored number of steps performed.
+  size_t stepsPerformed;
 };
 
 } // namespace rl

src/mlpack/methods/reinforcement_learning/environment/continuous_mountain_car.hpp

@@ -83,7 +83,7 @@ class ContinuousMountainCar
    */
   struct Action
   {
-    double action[1];
+    double action = 0.0;
     // Storing degree of freedom
     const int size = 1;
   };
@@ -97,19 +97,27 @@ class ContinuousMountainCar
    * @param velocityMin Minimum legal velocity.
    * @param velocityMax Maximum legal velocity.
    * @param power Power generated by car.
+   * @param doneReward Reward recieved by the agent on success.
+   * @param maxSteps The number of steps after which the episode
+   *    terminates. If the value is 0, there is no limit.
    */
   ContinuousMountainCar(const double positionMin = -1.2,
                         const double positionMax = 0.6,
                         const double positionGoal = 0.45,
                         const double velocityMin = -0.07,
                         const double velocityMax = 0.07,
-                        const double power = 0.0015) :
+                        const double power = 0.0015,
+                        const double doneReward = 100,
+                        const size_t maxSteps = 0) :
       positionMin(positionMin),
       positionMax(positionMax),
       positionGoal(positionGoal),
       velocityMin(velocityMin),
       velocityMax(velocityMax),
-      power(power)
+      power(power),
+      doneReward(doneReward),
+      maxSteps(maxSteps),
+      stepsPerformed(0)
   { /* Nothing to do here */ }
 
   /**
@@ -119,14 +127,16 @@ class ContinuousMountainCar
    * @param state The current state.
    * @param action The current action.
    * @param nextState The next state.
-   * @return reward, it's always -1.0.
    */
   double Sample(const State& state,
                 const Action& action,
-                State& nextState) const
+                State& nextState)
   {
+    // Update the number of steps performed.
+    stepsPerformed++;
+
     // Calculate acceleration.
-    double force = std::min(std::max(action.action[0], -1.0), 1.0);
+    double force = std::min(std::max(action.action, -1.0), 1.0);
 
     // Update states.
     nextState.Velocity() = state.Velocity() + force * power - 0.0025 *
@@ -139,13 +149,16 @@ class ContinuousMountainCar
     if (nextState.Position() == positionMin && nextState.Velocity() < 0)
       nextState.Velocity() = 0.0;
 
-    // Calculate reward
-    double reward = 0.0;
-    // If it is a terminal state, add a reward of 100.0
-    if (IsTerminal(nextState))
-      reward = 100.0;
-    reward -= std::pow(action.action[0], 2) * 0.1;
-    return reward;
+    // Check if the episode has terminated.
+    bool done = IsTerminal(nextState);
+
+    // Do not reward the agent if time ran out.
+    if (done && maxSteps != 0 && stepsPerformed >= maxSteps)
+      return 0;
+    else if (done)
+      return doneReward;
+
+    return std::pow(action.action, 2) * 0.1;
   }
 
   /**
@@ -156,7 +169,7 @@ class ContinuousMountainCar
    * @param action The current action.
    * @return reward, it's always -1.0.
    */
-  double Sample(const State& state, const Action& action) const
+  double Sample(const State& state, const Action& action)
   {
     State nextState;
     return Sample(state, action, nextState);
@@ -168,9 +181,10 @@ class ContinuousMountainCar
    *
    * @return Initial state for each episode.
    */
-  State InitialSample() const
+  State InitialSample()
   {
     State state;
+    stepsPerformed = 0;
     state.Velocity() = 0.0;
     state.Position() = math::Random(-0.6, -0.4);
     return state;
@@ -184,9 +198,28 @@ class ContinuousMountainCar
    */
   bool IsTerminal(const State& state) const
   {
-    return state.Position() >= positionGoal;
+    if (maxSteps != 0 && stepsPerformed >= maxSteps)
+    {
+      Log::Info << "Episode terminated due to the maximum number of steps"
+          "being taken.";
+      return true;
+    }
+    else if (state.Position() >= positionGoal)
+    {
+      Log::Info << "Episode terminated due to agent succeeding.";
+      return true;
+    }
+    return false;
   }
 
+  //! Get the number of steps performed.
+  size_t StepsPerformed() const { return stepsPerformed; }
+
+  //! Get the maximum number of steps allowed.
+  size_t MaxSteps() const { return maxSteps; }
+  //! Set the maximum number of steps allowed.
+  size_t& MaxSteps() { return maxSteps; }
+
  private:
   //! Locally-stored minimum legal position.
   double positionMin;
@@ -205,6 +238,15 @@ class ContinuousMountainCar
 
   //! Locally-stored power.
   double power;
+
+  //! Locally-stored done reward.
+  double doneReward;
+
+  //! Locally-stored maximum number of steps.
+  size_t maxSteps;
+
+  //! Locally-stored number of steps performed.
+  size_t stepsPerformed;
 };
 
 } // namespace rl

src/mlpack/methods/reinforcement_learning/environment/continuous_multiple_pole_cart.hpp

@@ -0,0 +1,310 @@
+/**
+ * @file continuous_multiple_pole_cart.hpp
+ * @author Rahul Ganesh Prabhu
+ *
+ * This file is an implementation of Continuous Multiple Pole Cart Balancing
+ * Task.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+
+#ifndef MLPACK_METHODS_RL_ENVIRONMENT_CONTINUOUS_MULTIPLE_POLE_CART_HPP
+#define MLPACK_METHODS_RL_ENVIRONMENT_CONTINUOUS_MULTIPLE_POLE_CART_HPP
+
+#include <mlpack/prereqs.hpp>
+
+namespace mlpack {
+namespace rl {
+
+/**
+ * Implementation of Continuous Multiple Pole Cart Balancing task.
+ */
+class ContinuousMultiplePoleCart
+{
+ public:
+  /**
+   * Implementation of the state of Continuous Multiple Pole Cart. The state is expressed as
+   * a matrix where the $0^{th}$ column is the state of the cart, represented by a tuple
+   * (position, velocity) and the $i^{th}$ column is the state of the $i^{th}$ pole, represented
+   * by a tuple (angle, angular velocity).
+   */
+  class State
+  {
+   public:
+    /**
+     * Construct a state instance.
+     * 
+     * @param numPoles The number of poles.
+     */
+    State(const size_t numPoles)
+    {
+      data = arma::zeros<arma::mat>(dimension, numPoles + 1);
+    }
+
+    /**
+     * Construct a state instance from given data.
+     *
+     * @param data Data for the position, velocity, angle and angular velocity.
+     */
+    State(const arma::mat& data) : data(data)
+    { /* Nothing to do here */ }
+
+    //! Modify the internal representation of the state.
+    arma::mat& Data() { return data; }
+
+    //! Get the position of the cart.
+    double Position() const { return data(0, 0); }
+    //! Modify the position of the cart.
+    double& Position() { return data(0, 0); }
+
+    //! Get the velocity of the cart.
+    double Velocity() const { return data(1, 0); }
+    //! Modify the velocity of the cart.
+    double& Velocity() { return data(1, 0); }
+
+    //! Get the angle of the $i^{th}$ pole with the vertical.
+    double Angle(const size_t i) const { return data(0, i); }
+    //! Modify the angle of the $i^{th}$ pole with the vertical.
+    double& Angle(const size_t i) { return data(0, i); }
+
+    //! Get the angular velocity of the $i^{th}$ pole.
+    double AngularVelocity(const size_t i) const { return data(1, i); }
+    //! Modify the angular velocity of the $i^{th}$ pole.
+    double& AngularVelocity(const size_t i) { return data(1, i); }
+
+    //! Encode the state to a matrix.
+    const arma::mat& Encode() const { return data; }
+
+    //! Dimension of the encoded state.
+    const size_t dimension = 2;
+
+   private:
+    //! Locally-stored state data.
+    arma::mat data;
+  };
+
+  /**
+   * Implementation of action of Continuous Multiple Pole Cart.
+   */
+  struct Action
+  {
+    double action = 0.0;
+    // Track the size of the action space.
+    const int size = 1;
+  };
+
+  /**
+   * Construct a Multiple Pole Cart instance using the given constants.
+   *
+   * @param poleNum The number of poles
+   * @param gravity The gravity constant.
+   * @param massCart The mass of the cart.
+   * @param massPole The mass of the pole.
+   * @param length The length of the pole.
+   * @param tau The time interval.
+   * @param thetaThresholdRadians The maximum angle.
+   * @param xThreshold The maximum position.
+   * @param doneReward Reward recieved by agent on success.
+   * @param maxSteps The number of steps after which the episode
+   *    terminates. If the value is 0, there is no limit.
+   */
+  ContinuousMultiplePoleCart(const size_t poleNum,
+                             const arma::vec& poleLengths,
+                             const arma::vec& poleMasses,
+                             const double gravity = 9.8,
+                             const double massCart = 1.0,
+                             const double tau = 0.02,
+                             const double thetaThresholdRadians = 12 * 2 *
+                                3.1416 / 360,
+                             const double xThreshold = 2.4,
+                             const double doneReward = 0.0,
+                             const size_t maxSteps = 0) :
+      poleNum(poleNum),
+      poleLengths(poleLengths),
+      poleMasses(poleMasses),
+      gravity(gravity),
+      massCart(massCart),
+      tau(tau),
+      thetaThresholdRadians(thetaThresholdRadians),
+      xThreshold(xThreshold),
+      doneReward(doneReward),
+      maxSteps(maxSteps),
+      stepsPerformed(0)
+  {
+    if (poleNum != poleLengths.n_elem)
+    {
+      Log::Fatal << "The number of lengths should be the same as the number of"
+          "poles." << std::endl;
+    }
+    if (poleNum != poleMasses.n_elem)
+    {
+      Log::Fatal << "The number of masses should be the same as the number of"
+          "poles." << std::endl;
+    }
+  }
+
+  /**
+   * Dynamics of Continuous Multiple Pole Cart instance. Get reward and next state 
+   * based on current state and current action.
+   *
+   * @param state The current state.
+   * @param action The current action.
+   * @param nextState The next state.
+   * @return reward, it's always 1.0.
+   */
+  double Sample(const State& state,
+                const Action& action,
+                State& nextState)
+  {
+    // Update the number of steps performed.
+    stepsPerformed++;
+
+    // Calculate acceleration.
+    double totalForce = action.action;
+    double totalMass = massCart;
+    for (size_t i = 0; i < poleNum; i++)
+    {
+      double poleOmega = state.AngularVelocity(i + 1);
+      double sinTheta = sin(state.Angle(i + 1));
+      totalForce += (poleMasses[i] * poleLengths[i] * poleOmega * poleOmega *
+          sinTheta) + 0.75 * poleMasses[i] * gravity * sin(2 * state.Angle(i +
+          1)) / 2;
+      totalMass += poleMasses[i] * (0.25 + 0.75 * sinTheta * sinTheta);
+    }
+    double xAcc = totalForce / totalMass;
+
+    // Update states of the poles.
+    for (size_t i = 1; i <= poleNum; i++)
+    {
+      double sinTheta = sin(state.Angle(i));
+      double cosTheta = cos(state.Angle(i));
+      nextState.Angle(i) = state.Angle(i) + tau * state.AngularVelocity(i);
+      nextState.AngularVelocity(i) = state.AngularVelocity(i) - tau * 0.75 *
+          (xAcc * cosTheta + gravity * sinTheta) / poleLengths[i - 1];
+    }
+
+    // Update state of the cart.
+    nextState.Position() = state.Position() + tau * state.Velocity();
+    nextState.Velocity() = state.Velocity() + tau * xAcc;
+
+    // Check if the episode has terminated.
+    bool done = IsTerminal(nextState);
+
+    // Do not reward agent if it failed.
+    if (done && maxSteps != 0 && stepsPerformed >= maxSteps)
+      return doneReward;
+    else if (done)
+      return 0;
+
+    /**
+     * When done is false, it means that the cartpole has fallen down.
+     * For this case the reward is 1.0.
+     */
+    return 1.0;
+  }
+
+  /**
+   * Dynamics of Continuous Multiple Pole Cart. Get reward based on current
+   * state and current action.
+   *
+   * @param state The current state.
+   * @param action The current action.
+   * @return reward, it's always 1.0.
+   */
+  double Sample(const State& state, const Action& action)
+  {
+    State nextState(poleNum);
+    return Sample(state, action, nextState);
+  }
+
+  /**
+   * Initial state representation is randomly generated within [-0.05, 0.05].
+   *
+   * @return Initial state for each episode.
+   */
+  State InitialSample()
+  {
+    stepsPerformed = 0;
+    return State((arma::randu<arma::mat>(2, poleNum + 1) - 0.5) / 10.0);
+  }
+
+  /**
+   * This function checks if the cart has reached the terminal state.
+   *
+   * @param state The desired state.
+   * @return true if state is a terminal state, otherwise false.
+   */
+  bool IsTerminal(const State& state) const
+  {
+    if (maxSteps != 0 && stepsPerformed >= maxSteps)
+    {
+      Log::Info << "Episode terminated due to the maximum number of steps"
+          "being taken.";
+      return true;
+    }
+    if (std::abs(state.Position()) > xThreshold)
+    {
+      Log::Info << "Episode terminated due to cart crossing threshold";
+      return true;
+    }
+    for (size_t i = 1; i <= poleNum; i++)
+    {
+      if (std::abs(state.Angle(i)) > thetaThresholdRadians)
+      {
+        Log::Info << "Episode terminated due to pole falling";
+        return true;
+      }
+    }
+    return false;
+  }
+
+  //! Get the number of steps performed.
+  size_t StepsPerformed() const { return stepsPerformed; }
+
+  //! Get the maximum number of steps allowed.
+  size_t MaxSteps() const { return maxSteps; }
+  //! Set the maximum number of steps allowed.
+  size_t& MaxSteps() { return maxSteps; }
+
+ private:
+  //! Locally-stored number of poles.
+  size_t poleNum;
+
+  //! Locally-stored length of poles.
+  arma::vec poleLengths;
+
+  //! Locally-stored mass of the pole.
+  arma::vec poleMasses;
+
+  //! Locally-stored gravity.
+  double gravity;
+
+  //! Locally-stored mass of the cart.
+  double massCart;
+
+  //! Locally-stored time interval.
+  double tau;
+
+  //! Locally-stored maximum angle.
+  double thetaThresholdRadians;
+
+  //! Locally-stored maximum position.
+  double xThreshold;
+
+  //! Locally-stored done reward.
+  double doneReward;
+
+  //! Locally-stored maximum number of steps.
+  size_t maxSteps;
+
+  //! Locally-stored number of steps performed.
+  size_t stepsPerformed;
+};
+
+} // namespace rl
+} // namespace mlpack
+
+#endif

src/mlpack/methods/reinforcement_learning/environment/mountain_car.hpp

@@ -43,7 +43,7 @@ class MountainCar
     /**
      * Construct a state based on the given data.
      *
-     * @param data Data for the velocityand position.
+     * @param data Data for the velocity and position.
      */
     State(const arma::colvec& data): data(data)
     { /* Nothing to do here. */ }
@@ -93,19 +93,25 @@ class MountainCar
    * @param positionGoal Final target position.
    * @param velocityMin Minimum legal velocity.
    * @param velocityMax Maximum legal velocity.
+   * @param doneReward The reward recieved by the agent on success.
+   * @param maxSteps The number of steps after which the episode
+   *    terminates. If the value is 0, there is no limit.
    */
   MountainCar(const double positionMin = -1.2,
               const double positionMax = 0.6,
               const double positionGoal = 0.5,
               const double velocityMin = -0.07,
               const double velocityMax = 0.07,
-              const double doneReward = 0) :
+              const double doneReward = 0,
+              const size_t maxSteps = 0) :
       positionMin(positionMin),
       positionMax(positionMax),
       positionGoal(positionGoal),
       velocityMin(velocityMin),
       velocityMax(velocityMax),
-      doneReward(doneReward)
+      doneReward(doneReward),
+      maxSteps(maxSteps),
+      stepsPerformed(0)
   { /* Nothing to do here */ }
 
   /**
@@ -119,8 +125,11 @@ class MountainCar
    */
   double Sample(const State& state,
                 const Action& action,
-                State& nextState) const
+                State& nextState)
   {
+    // Update the number of steps performed.
+    stepsPerformed++;
+
     // Calculate acceleration.
     int direction = action - 1;
     nextState.Velocity() = state.Velocity() + 0.001 * direction - 0.0025 *
@@ -136,17 +145,16 @@ class MountainCar
     if (nextState.Position() == positionMin && nextState.Velocity() < 0)
       nextState.Velocity() = 0.0;
 
+    // Check if the episode has terminated.
     bool done = IsTerminal(nextState);
-    /**
-     * If done is true , it means that car has reached its goal.
-     * To make sure that the agent learns this, we will give some
-     * positive reward to the agent. If the agent doesn't reach the
-     * terminal state, then we will give a -1.0 reward to penalize 
-     * the agent to take that step.
-     */
-    if (done)
+
+    // Do not reward the agent if time ran out.
+    if (done && maxSteps != 0 && stepsPerformed >= maxSteps)
+      return 0;
+    else if (done)
       return doneReward;
-    return -1.0;
+
+    return -1;
   }
 
   /**
@@ -157,7 +165,7 @@ class MountainCar
    * @param action The current action.
    * @return reward, it's always -1.0.
    */
-  double Sample(const State& state, const Action& action) const
+  double Sample(const State& state, const Action& action)
   {
     State nextState;
     return Sample(state, action, nextState);
@@ -169,25 +177,45 @@ class MountainCar
    *
    * @return Initial state for each episode.
    */
-  State InitialSample() const
+  State InitialSample()
   {
     State state;
+    stepsPerformed = 0;
     state.Velocity() = 0.0;
     state.Position() = arma::as_scalar(arma::randu(1)) * 0.2 - 0.6;
     return state;
   }
 
   /**
-   * Whether given state is a terminal state.
+   * This function checks if the car has reached the terminal state.
    *
    * @param state desired state.
    * @return true if state is a terminal state, otherwise false.
    */
   bool IsTerminal(const State& state) const
   {
-    return state.Position() >= positionGoal;
+    if (maxSteps != 0 && stepsPerformed >= maxSteps)
+    {
+      Log::Info << "Episode terminated due to the maximum number of steps"
+          "being taken.";
+      return true;
+    }
+    else if (state.Position() >= positionGoal)
+    {
+      Log::Info << "Episode terminated due to agent succeeding.";
+      return true;
+    }
+    return false;
   }
 
+  //! Get the number of steps performed.
+  size_t StepsPerformed() const { return stepsPerformed; }
+
+  //! Get the maximum number of steps allowed.
+  size_t MaxSteps() const { return maxSteps; }
+  //! Set the maximum number of steps allowed.
+  size_t& MaxSteps() { return maxSteps; }
+
  private:
   //! Locally-stored minimum legal position.
   double positionMin;
@@ -206,6 +234,12 @@ class MountainCar
 
   //! Locally-stored done reward.
   double doneReward;
+
+  //! Locally-stored maximum number of steps.
+  size_t maxSteps;
+
+  //! Locally-stored number of steps performed.
+  size_t stepsPerformed;
 };
 
 } // namespace rl

src/mlpack/methods/reinforcement_learning/environment/multiple_pole_cart.hpp

@@ -0,0 +1,316 @@
+/**
+ * @file multiple_pole_cart.hpp
+ * @author Rahul Ganesh Prabhu
+ *
+ * This file is an implementation of Multiple Pole Cart Balancing Task.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+
+#ifndef MLPACK_METHODS_RL_ENVIRONMENT_MULTIPLE_POLE_CART_HPP
+#define MLPACK_METHODS_RL_ENVIRONMENT_MULTIPLE_POLE_CART_HPP
+
+#include <mlpack/prereqs.hpp>
+
+namespace mlpack {
+namespace rl {
+
+/**
+ * Implementation of Multiple Pole Cart Balancing task.
+ */
+class MultiplePoleCart
+{
+ public:
+  /**
+   * Implementation of the state of Multiple Pole Cart. The state is expressed as
+   * a matrix where the $0^{th}$ column is the state of the cart, represented by a tuple
+   * (position, velocity) and the $i^{th}$ column is the state of the $i^{th}$ pole, represented
+   * by a tuple (angle, angular velocity).
+   */
+  class State
+  {
+   public:
+    /**
+     * Construct a state instance.
+     * 
+     * @param numPoles The number of poles.
+     */
+    State(const size_t numPoles)
+    {
+      data = arma::zeros<arma::mat>(dimension, numPoles + 1);
+    }
+
+    /**
+     * Construct a state instance from given data.
+     *
+     * @param data Data for the position, velocity, angle and angular velocity.
+     */
+    State(const arma::mat& data) : data(data)
+    { /* Nothing to do here */ }
+
+    //! Modify the internal representation of the state.
+    arma::mat& Data() { return data; }
+
+    //! Get the position of the cart.
+    double Position() const { return data(0, 0); }
+    //! Modify the position of the cart.
+    double& Position() { return data(0, 0); }
+
+    //! Get the velocity of the cart.
+    double Velocity() const { return data(1, 0); }
+    //! Modify the velocity of the cart.
+    double& Velocity() { return data(1, 0); }
+
+    //! Get the angle of the $i^{th}$ pole with the vertical.
+    double Angle(const size_t i) const { return data(0, i); }
+    //! Modify the angle of the $i^{th}$ pole with the vertical.
+    double& Angle(const size_t i) { return data(0, i); }
+
+    //! Get the angular velocity of the $i^{th}$ pole.
+    double AngularVelocity(const size_t i) const { return data(1, i); }
+    //! Modify the angular velocity of the $i^{th}$ pole.
+    double& AngularVelocity(const size_t i) { return data(1, i); }
+
+    //! Encode the state to a matrix.
+    const arma::mat& Encode() const { return data; }
+
+    //! Dimension of the encoded state.
+    const size_t dimension = 2;
+
+   private:
+    //! Locally-stored state data.
+    arma::mat data;
+  };
+
+  /**
+   * Implementation of action of Multiple Pole Cart.
+   */
+  enum Action
+  {
+    backward,
+    forward,
+
+    // Track the size of the action space.
+    size
+  };
+
+  /**
+   * Construct a Multiple Pole Cart instance using the given constants.
+   *
+   * @param poleNum The number of poles
+   * @param gravity The gravity constant.
+   * @param massCart The mass of the cart.
+   * @param massPole The mass of the pole.
+   * @param length The length of the pole.
+   * @param forceMag The magnitude of the applied force.
+   * @param tau The time interval.
+   * @param thetaThresholdRadians The maximum angle.
+   * @param xThreshold The maximum position.
+   * @param doneReward Reward recieved by agent on success.
+   * @param maxSteps The number of steps after which the episode
+   *    terminates. If the value is 0, there is no limit.
+   */
+  MultiplePoleCart(const size_t poleNum,
+                   const arma::vec& poleLengths,
+                   const arma::vec& poleMasses,
+                   const double gravity = 9.8,
+                   const double massCart = 1.0,
+                   const double forceMag = 10.0,
+                   const double tau = 0.02,
+                   const double thetaThresholdRadians = 12 * 2 * 3.1416 / 360,
+                   const double xThreshold = 2.4,
+                   const double doneReward = 0.0,
+                   const size_t maxSteps = 0) :
+      poleNum(poleNum),
+      poleLengths(poleLengths),
+      poleMasses(poleMasses),
+      gravity(gravity),
+      massCart(massCart),
+      forceMag(forceMag),
+      tau(tau),
+      thetaThresholdRadians(thetaThresholdRadians),
+      xThreshold(xThreshold),
+      doneReward(doneReward),
+      maxSteps(maxSteps),
+      stepsPerformed(0)
+  {
+    if (poleNum != poleLengths.n_elem)
+    {
+      Log::Fatal << "The number of lengths should be the same as the number of"
+          "poles." << std::endl;
+    }
+    if (poleNum != poleMasses.n_elem)
+    {
+      Log::Fatal << "The number of masses should be the same as the number of"
+          "poles." << std::endl;
+    }
+  }
+
+  /**
+   * Dynamics of Multiple Pole Cart instance. Get reward and next state based on current
+   * state and current action.
+   *
+   * @param state The current state.
+   * @param action The current action.
+   * @param nextState The next state.
+   * @return reward, it's always 1.0.
+   */
+  double Sample(const State& state,
+                const Action& action,
+                State& nextState)
+  {
+    // Update the number of steps performed.
+    stepsPerformed++;
+
+    // Calculate acceleration.
+    double totalForce = action ? forceMag : -forceMag;
+    double totalMass = massCart;
+    for (size_t i = 0; i < poleNum; i++)
+    {
+      double poleOmega = state.AngularVelocity(i + 1);
+      double sinTheta = sin(state.Angle(i + 1));
+      totalForce += (poleMasses[i] * poleLengths[i] * poleOmega * poleOmega *
+          sinTheta) + 0.75 * poleMasses[i] * gravity * sin(2 * state.Angle(i +
+          1)) / 2;
+      totalMass += poleMasses[i] * (0.25 + 0.75 * sinTheta * sinTheta);
+    }
+    double xAcc = totalForce / totalMass;
+
+    // Update states of the poles.
+    for (size_t i = 1; i <= poleNum; i++)
+    {
+      double sinTheta = sin(state.Angle(i));
+      double cosTheta = cos(state.Angle(i));
+      nextState.Angle(i) = state.Angle(i) + tau * state.AngularVelocity(i);
+      nextState.AngularVelocity(i) = state.AngularVelocity(i) - tau * 0.75 *
+          (xAcc * cosTheta + gravity * sinTheta) / poleLengths[i - 1];
+    }
+
+    // Update state of the cart.
+    nextState.Position() = state.Position() + tau * state.Velocity();
+    nextState.Velocity() = state.Velocity() + tau * xAcc;
+
+    // Check if the episode has terminated.
+    bool done = IsTerminal(nextState);
+
+    // Do not reward agent if it failed.
+    if (done && maxSteps != 0 && stepsPerformed >= maxSteps)
+      return doneReward;
+    else if (done)
+      return 0;
+
+    /**
+     * When done is false, it means that the cartpole has fallen down.
+     * For this case the reward is 1.0.
+     */
+    return 1.0;
+  }
+
+  /**
+   * Dynamics of Multiple Pole Cart. Get reward based on current state and current
+   * action.
+   *
+   * @param state The current state.
+   * @param action The current action.
+   * @return reward, it's always 1.0.
+   */
+  double Sample(const State& state, const Action& action)
+  {
+    State nextState(poleNum);
+    return Sample(state, action, nextState);
+  }
+
+  /**
+   * Initial state representation is randomly generated within [-0.05, 0.05].
+   *
+   * @return Initial state for each episode.
+   */
+  State InitialSample()
+  {
+    stepsPerformed = 0;
+    return State((arma::randu<arma::mat>(2, poleNum + 1) - 0.5) / 10.0);
+  }
+
+  /**
+   * This function checks if the car has reached the terminal state.
+   *
+   * @param state The desired state.
+   * @return true if state is a terminal state, otherwise false.
+   */
+  bool IsTerminal(const State& state) const
+  {
+    if (maxSteps != 0 && stepsPerformed >= maxSteps)
+    {
+      Log::Info << "Episode terminated due to the maximum number of steps"
+          "being taken.";
+      return true;
+    }
+    if (std::abs(state.Position()) > xThreshold)
+    {
+      Log::Info << "Episode terminated due to cart crossing threshold";
+      return true;
+    }
+    for (size_t i = 1; i <= poleNum; i++)
+    {
+      if (std::abs(state.Angle(i)) > thetaThresholdRadians)
+      {
+        Log::Info << "Episode terminated due to pole falling";
+        return true;
+      }
+    }
+    return false;
+  }
+
+  //! Get the number of steps performed.
+  size_t StepsPerformed() const { return stepsPerformed; }
+
+  //! Get the maximum number of steps allowed.
+  size_t MaxSteps() const { return maxSteps; }
+  //! Set the maximum number of steps allowed.
+  size_t& MaxSteps() { return maxSteps; }
+
+ private:
+  //! Locally-stored number of poles.
+  size_t poleNum;
+
+  //! Locally-stored length of poles.
+  arma::vec poleLengths;
+
+  //! Locally-stored mass of the pole.
+  arma::vec poleMasses;
+
+  //! Locally-stored gravity.
+  double gravity;
+
+  //! Locally-stored mass of the cart.
+  double massCart;
+
+  //! Locally-stored magnitude of the applied force.
+  double forceMag;
+
+  //! Locally-stored time interval.
+  double tau;
+
+  //! Locally-stored maximum angle.
+  double thetaThresholdRadians;
+
+  //! Locally-stored maximum position.
+  double xThreshold;
+
+  //! Locally-stored done reward.
+  double doneReward;
+
+  //! Locally-stored maximum number of steps.
+  size_t maxSteps;
+
+  //! Locally-stored number of steps performed.
+  size_t stepsPerformed;
+};
+
+} // namespace rl
+} // namespace mlpack
+
+#endif

src/mlpack/methods/reinforcement_learning/environment/pendulum.hpp

@@ -82,7 +82,7 @@ class Pendulum
    */
   struct Action
   {
-    double action[1];
+    double action = 0.0;
     // Storing degree of freedom
     const int size = 1;
   };
@@ -93,13 +93,25 @@ class Pendulum
    * @param maxAngularVelocity Maximum angular velocity.
    * @param maxTorque Maximum torque.
    * @param dt The differential value.
+   * @param angleThreshold The region about the upright position where the
+   *    state is considered terminal.
+   * @param doneReward The reward recieved by the agent on success.
+   * @param maxSteps The number of steps after which the episode
+   *    terminates. If the value is 0, there is no limit.
    */
   Pendulum(const double maxAngularVelocity = 8,
            const double maxTorque = 2.0,
-           const double dt = 0.05) :
+           const double dt = 0.05,
+           const double angleThreshold = M_PI / 12,
+           const double doneReward = 0.0,
+           const size_t maxSteps = 0) :
       maxAngularVelocity(maxAngularVelocity),
       maxTorque(maxTorque),
-      dt(dt)
+      dt(dt),
+      angleThreshold(angleThreshold),
+      doneReward(doneReward),
+      maxSteps(maxSteps),
+      stepsPerformed(0)
   { /* Nothing to do here */ }
 
   /**
@@ -113,8 +125,11 @@ class Pendulum
    */
   double Sample(const State& state,
                 const Action& action,
-                State& nextState) const
+                State& nextState)
   {
+    // Update the number of steps performed.
+    stepsPerformed++;
+
     // Get current state.
     double theta = state.Theta();
     double angularVelocity = state.AngularVelocity();
@@ -126,7 +141,7 @@ class Pendulum
 
     // Get action and clip the values between max and min limits.
     double torque = std::min(
-        std::max(action.action[0], -maxTorque), maxTorque);
+        std::max(action.action, -maxTorque), maxTorque);
 
     // Calculate costs of taking this action in the current state.
     double costs = std::pow(AngleNormalize(theta), 2) + 0.1 *
@@ -140,6 +155,15 @@ class Pendulum
         -maxAngularVelocity), maxAngularVelocity);
     nextState.Theta() = theta + newAngularVelocity * dt;
 
+    // Check if the episode has terminated
+    bool done = IsTerminal(nextState);
+
+    // Do not reward the agent if time ran out.
+    if (done && maxSteps != 0 && stepsPerformed >= maxSteps)
+      return 0;
+    else if (done)
+      return doneReward;
+
     // Return the reward of taking the action in current state.
     // The reward is simply the negative of cost incurred for the action.
     return -costs;
@@ -152,7 +176,7 @@ class Pendulum
    * @param action The current action.
    * @return reward, The reward.
    */
-  double Sample(const State& state, const Action& action) const
+  double Sample(const State& state, const Action& action)
   {
     State nextState;
     return Sample(state, action, nextState);
@@ -164,16 +188,17 @@ class Pendulum
    *
    * @return Initial state for each episode.
    */
-  State InitialSample() const
+  State InitialSample()
   {
     State state;
-    state.Theta() = math::Random(-M_PI, M_PI);
+    state.Theta() = math::Random(-M_PI + angleThreshold, M_PI - angleThreshold);
     state.AngularVelocity() = math::Random(-1.0, 1.0);
+    stepsPerformed = 0;
     return state;
   }
 
   /**
-   * This function calculates the normalized anlge for a particular theta.
+   * This function calculates the normalized angle for a particular theta.
    *
    * @param theta The un-normalized angle.
    */
@@ -183,6 +208,37 @@ class Pendulum
     return double(fmod(theta + M_PI, 2 * M_PI) - M_PI);
   }
 
+  /**
+   * This function checks if the pendulum has reaches a terminal state
+   * 
+   * @param state desired state.
+   * @return true if state is a terminal state, otherwise false.
+   */
+  bool IsTerminal(const State& state) const
+  {
+    if (maxSteps != 0 && stepsPerformed >= maxSteps)
+    {
+      Log::Info << "Episode terminated due to the maximum number of steps"
+          "being taken.";
+      return true;
+    }
+    else if (state.Theta() > M_PI - angleThreshold ||
+        state.Theta() < -M_PI + angleThreshold)
+    {
+      Log::Info << "Episode terminated due to agent succeeding.";
+      return true;
+    }
+    return false;
+  }
+
+  //! Get the number of steps performed.
+  size_t StepsPerformed() const { return stepsPerformed; }
+
+  //! Get the maximum number of steps allowed.
+  size_t MaxSteps() const { return maxSteps; }
+  //! Set the maximum number of steps allowed.
+  size_t& MaxSteps() { return maxSteps; }
+
  private:
   //! Locally-stored maximum legal angular velocity.
   double maxAngularVelocity;
@@ -192,6 +248,18 @@ class Pendulum
 
   //! Locally-stored dt.
   double dt;
+
+  //! Locally-stored angle threshold.
+  double angleThreshold;
+
+  //! Locally-stored done reward.
+  double doneReward;
+
+  //! Locally-stored maximum number of steps.
+  size_t maxSteps;
+
+  //! Locally-stored number of steps performed.
+  size_t stepsPerformed;
 };
 
 } // namespace rl

src/mlpack/methods/reinforcement_learning/environment/reward_clipping.hpp

@@ -86,7 +86,7 @@ class RewardClipping
    */
   double Sample(const State& state,
                 const Action& action,
-                State& nextState) const
+                State& nextState)
   {
     // Get original unclipped reward from base environment.
     double unclippedReward =  environment.Sample(state, action, nextState);
@@ -102,7 +102,7 @@ class RewardClipping
    * @param action The current action.
    * @return clippedReward, Reward clipped between [minReward, maxReward].
    */
-  double Sample(const State& state, const Action& action) const
+  double Sample(const State& state, const Action& action)
   {
     State nextState;
     return Sample(state, action, nextState);

src/mlpack/methods/reinforcement_learning/q_learning.hpp

@@ -16,6 +16,7 @@
 #include <mlpack/prereqs.hpp>
 
 #include "replay/random_replay.hpp"
+#include "replay/prioritized_replay.hpp"
 #include "training_config.hpp"
 
 namespace mlpack {

src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp

@@ -48,6 +48,7 @@ QLearning<
   // Set up q-learning network.
   if (learningNetwork.Parameters().is_empty())
     learningNetwork.ResetParameters();
+
   this->updater.Initialize(learningNetwork.Parameters().n_rows,
       learningNetwork.Parameters().n_cols);
   targetNetwork = learningNetwork;
@@ -123,6 +124,7 @@ double QLearning<
   arma::colvec sampledRewards;
   arma::mat sampledNextStates;
   arma::icolvec isTerminal;
+
   replayMethod.Sample(sampledStates, sampledActions, sampledRewards,
       sampledNextStates, isTerminal);
 
@@ -154,15 +156,23 @@ double QLearning<
   for (size_t i = 0; i < sampledNextStates.n_cols; ++i)
   {
     if (isTerminal[i])
-      target(sampledActions[i], i) = sampledRewards[i];
+    {
+      target(sampledActions(i), i) = sampledRewards(i);
+    }
     else
-      target(sampledActions[i], i) = sampledRewards[i] + config.Discount() *
-          nextActionValues(bestActions[i], i);
+    {
+      target(sampledActions(i), i) = sampledRewards(i) + config.Discount() *
+          nextActionValues(bestActions(i), i);
+    }
   }
 
   // Learn form experience.
   arma::mat gradients;
   learningNetwork.Backward(target, gradients);
+
+  replayMethod.Update(target, sampledActions,
+      nextActionValues, gradients);
+
   updater.Update(learningNetwork.Parameters(), config.StepSize(), gradients);
 
   return reward;

src/mlpack/methods/reinforcement_learning/replay/CMakeLists.txt

@@ -2,6 +2,8 @@
 # Anything not in this list will not be compiled into mlpack.
 set(SOURCES
   random_replay.hpp
+  sumtree.hpp
+  prioritized_replay.hpp
 )
 
 # Add directory name to sources.

src/mlpack/methods/reinforcement_learning/replay/prioritized_replay.hpp

@@ -0,0 +1,292 @@
+/**
+ * @file prioritized_experience_replay.hpp
+ * @author Xiaohong
+ *
+ * This file is an implementation of prioritized experience replay.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_METHODS_RL_PRIORITIZED_REPLAY_HPP
+#define MLPACK_METHODS_RL_PRIORITIZED_REPLAY_HPP
+
+#include <mlpack/prereqs.hpp>
+#include "sumtree.hpp"
+
+namespace mlpack {
+namespace rl {
+
+/**
+ * Implementation of prioritized experience replay. Prioritized experience
+ * replay can replay important transitions more frequently by prioritizing
+ * transitions, and make agent learn more efficiently.
+ *
+ * @code
+ * @article{schaul2015prioritized,
+ *  title   = {Prioritized experience replay},
+ *  author  = {Schaul, Tom and Quan, John and Antonoglou,
+ *             Ioannis and Silver, David},
+ *  journal = {arXiv preprint arXiv:1511.05952},
+ *  year    = {2015}
+ *  }
+ * @endcode
+ *
+ * @tparam EnvironmentType Desired task.
+ */
+template <typename EnvironmentType>
+class PrioritizedReplay
+{
+ public:
+  //! Convenient typedef for action.
+  using ActionType = typename EnvironmentType::Action;
+
+  //! Convenient typedef for state.
+  using StateType = typename EnvironmentType::State;
+
+  /**
+   * Default constructor.
+   */
+  PrioritizedReplay()
+  { /* Nothing to do here. */ }
+
+  /**
+   * Construct an instance of prioritized experience replay class.
+   *
+   * @param batchSize Number of examples returned at each sample.
+   * @param capacity Total memory size in terms of number of examples.
+   * @param alpha How much prioritization is used.
+   * @param dimension The dimension of an encoded state.
+   */
+  PrioritizedReplay(const size_t batchSize,
+                    const size_t capacity,
+                    const double alpha,
+                    const size_t dimension = StateType::dimension) :
+      batchSize(batchSize),
+      capacity(capacity),
+      position(0),
+      states(dimension, capacity),
+      actions(capacity),
+      rewards(capacity),
+      nextStates(dimension, capacity),
+      isTerminal(capacity),
+      full(false),
+      alpha(alpha),
+      maxPriority(1.0),
+      initialBeta(0.6),
+      replayBetaIters(10000)
+  {
+    size_t size = 1;
+    while (size < capacity)
+    {
+      size *= 2;
+    }
+
+    beta = initialBeta;
+    idxSum = SumTree<double>(size);
+  }
+
+  /**
+   * Store the given experience and set the priorities for the given experience.
+   *
+   * @param state Given state.
+   * @param action Given action.
+   * @param reward Given reward.
+   * @param nextState Given next state.
+   * @param isEnd Whether next state is terminal state.
+   */
+  void Store(const StateType& state,
+             ActionType action,
+             double reward,
+             const StateType& nextState,
+             bool isEnd)
+  {
+    states.col(position) = state.Encode();
+    actions(position) = action;
+    rewards(position) = reward;
+    nextStates.col(position) = nextState.Encode();
+    isTerminal(position) = isEnd;
+
+    idxSum.Set(position, maxPriority * alpha);
+
+    position++;
+    if (position == capacity)
+    {
+      full = true;
+      position = 0;
+    }
+  }
+
+  /**
+   * Sample some experience according to their priorities.
+   *
+   * @return The indices to be chosen.
+   */
+  arma::ucolvec SampleProportional()
+  {
+    arma::ucolvec idxes(batchSize);
+    double totalSum = idxSum.Sum(0, (full ? capacity : position));
+    double sumPerRange = totalSum / batchSize;
+    for (size_t bt = 0; bt < batchSize; bt++)
+    {
+      const double mass = arma::randu() * sumPerRange + bt * sumPerRange;
+      idxes(bt) = idxSum.FindPrefixSum(mass);
+    }
+    return idxes;
+  }
+
+  /**
+   * Sample some experience according to their priorities.
+   *
+   * @param sampledStates Sampled encoded states.
+   * @param sampledActions Sampled actions.
+   * @param sampledRewards Sampled rewards.
+   * @param sampledNextStates Sampled encoded next states.
+   * @param isTerminal Indicate whether corresponding next state is terminal
+   *        state.
+   */
+  void Sample(arma::mat& sampledStates,
+              arma::icolvec& sampledActions,
+              arma::colvec& sampledRewards,
+              arma::mat& sampledNextStates,
+              arma::icolvec& isTerminal)
+  {
+    sampledIndices = SampleProportional();
+    BetaAnneal();
+
+    sampledStates = states.cols(sampledIndices);
+    sampledActions = actions.elem(sampledIndices);
+    sampledRewards = rewards.elem(sampledIndices);
+    sampledNextStates = nextStates.cols(sampledIndices);
+    isTerminal = this->isTerminal.elem(sampledIndices);
+
+    // Calculate the weights of sampled transitions.
+
+    size_t numSample = full ? capacity : position;
+    weights = arma::rowvec(sampledIndices.n_rows);
+
+    for (size_t i = 0; i < sampledIndices.n_rows; i++)
+    {
+      double p_sample = idxSum.Get(sampledIndices(i)) / idxSum.Sum();
+      weights(i) = pow(numSample * p_sample, -beta);
+    }
+    weights /= weights.max();
+  }
+
+  /**
+   * Update priorities of sampled transitions.
+   *
+   * @param indices The indices of sample to be updated.
+   * @param priorities Their corresponding priorities.
+   */
+  void UpdatePriorities(arma::ucolvec& indices, arma::colvec& priorities)
+  {
+      arma::colvec alphaPri = alpha * priorities;
+      maxPriority = std::max(maxPriority, arma::max(priorities));
+      idxSum.BatchUpdate(indices, alphaPri);
+  }
+
+  /**
+   * Get the number of transitions in the memory.
+   *
+   * @return Actual used memory size.
+   */
+  const size_t& Size()
+  {
+    return full ? capacity : position;
+  }
+
+  /**
+   * Annealing the beta.
+   */
+  void BetaAnneal()
+  {
+    beta = beta + (1 - initialBeta) * 1.0 / replayBetaIters;
+  }
+
+  /**
+   * Update the priorities of transitions and Update the gradients.
+   *
+   * @param target The learned value.
+   * @param sampledActions Agent's sampled action.
+   * @param nextActionValues Agent's next action.
+   * @param gradients The model's gradients.
+   */
+  void Update(arma::mat target,
+              arma::icolvec sampledActions,
+              arma::mat nextActionValues,
+              arma::mat& gradients)
+  {
+    arma::colvec tdError(target.n_cols);
+    for (size_t i = 0; i < target.n_cols; i ++)
+    {
+      tdError(i) = nextActionValues(sampledActions(i), i) -
+          target(sampledActions(i), i);
+    }
+    tdError = arma::abs(tdError);
+    UpdatePriorities(sampledIndices, tdError);
+
+    // Update the gradient
+    gradients = arma::mean(weights) * gradients;
+  }
+
+
+ private:
+  //! Locally-stored number of examples of each sample.
+  size_t batchSize;
+
+  //! Locally-stored total memory limit.
+  size_t capacity;
+
+  //! Indicate the position to store new transition.
+  size_t position;
+
+  //! Locally-stored encoded previous states.
+  arma::mat states;
+
+  //! Locally-stored previous actions.
+  arma::icolvec actions;
+
+  //! Locally-stored previous rewards.
+  arma::colvec rewards;
+
+  //! Locally-stored encoded previous next states.
+  arma::mat nextStates;
+
+  //! Locally-stored termination information of previous experience.
+  arma::icolvec isTerminal;
+
+  //! Locally-stored indicator that whether the memory is full or not.
+  bool full;
+
+  //! How much prioritization is used.
+  //! (0 - no prioritization, 1 - full prioritization)
+  double alpha;
+
+  //! Locally-stored the max priority.
+  double maxPriority;
+
+  //! Initial value of beta for prioritized replay buffer.
+  double initialBeta;
+
+  //! The value of beta for current sample.
+  double beta;
+
+  //! How many iteration for replay beta to decay.
+  size_t replayBetaIters;
+
+  //! Locally-stored the prefix sum of prioritization.
+  SumTree<double> idxSum;
+
+  //! Locally-stored the indices of sampled transitions.
+  arma::ucolvec sampledIndices;
+
+  //! Locally-stored the weights of sampled transitions.
+  arma::rowvec weights;
+};
+
+} // namespace rl
+} // namespace mlpack
+
+#endif

src/mlpack/methods/reinforcement_learning/replay/random_replay.hpp

@@ -49,6 +49,13 @@ class RandomReplay
   //! Convenient typedef for state.
   using StateType = typename EnvironmentType::State;
 
+  RandomReplay():
+      batchSize(0),
+      capacity(0),
+      position(0),
+      full(false)
+  { /* Nothing to do here. */ }
+
   /**
    * Construct an instance of random experience replay class.
    *
@@ -135,6 +142,22 @@ class RandomReplay
     return full ? capacity : position;
   }
 
+  /**
+   * Update the priorities of transitions and Update the gradients.
+   *
+   * @param target The learned value
+   * @param sampledActions Agent's sampled action
+   * @param nextActionValues Agent's next action
+   * @param gradients The model's gradients
+   */
+  void Update(arma::mat /* target */,
+              arma::icolvec /* sampledActions */,
+              arma::mat /* nextActionValues */,
+              arma::mat& /* gradients */)
+  {
+    /* Do nothing for random replay. */
+  }
+
  private:
   //! Locally-stored number of examples of each sample.
   size_t batchSize;

src/mlpack/methods/reinforcement_learning/replay/sumtree.hpp

@@ -0,0 +1,192 @@
+/**
+ * @file sumtree.hpp
+ * @author Xiaohong
+ *
+ * This file is an implementation of sumtree. Based on:
+ * https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_METHODS_RL_SUMTREE_HPP
+#define MLPACK_METHODS_RL_SUMTREE_HPP
+
+#include <mlpack/prereqs.hpp>
+
+namespace mlpack {
+namespace rl {
+
+/**
+ * Implementation of SumTree.
+ *
+ * Build a Segment Tree like data structure.
+ * https://en.wikipedia.org/wiki/Segment_tree
+ *
+ * Used to maintain prefix-sum of an array.
+ *
+ * @tparam T The array's element type.
+ */
+template<typename T>
+class SumTree
+{
+ public:
+  /**
+   * Default constructor.
+   */
+  SumTree() : capacity(0)
+  { /* Nothing to do here. */ }
+
+  /**
+   * Construct an instance of SumTree class.
+   *
+   * @param capacity Size of data.
+   */
+  SumTree(const size_t capacity) : capacity(capacity)
+  {
+    element = std::vector<T>(2 * capacity);
+  }
+
+  /**
+   * Set the data array with idx.
+   *
+   * @param idx The array idx to be changed.
+   * @param value The data that array with idx to be.
+   */
+  void Set(size_t idx, const T value)
+  {
+    idx += capacity;
+    element[idx] = value;
+    idx /= 2;
+    while (idx >= 1)
+    {
+      element[idx] = element[2 * idx] + element[2 * idx + 1];
+      idx /= 2;
+    }
+  }
+
+  /**
+   * Update the data with batch rather loop over the indices with set method.
+   *
+   * @param indices The indices of data to be changed.
+   * @param data The data that array with indices to be.
+   */
+  void BatchUpdate(const arma::ucolvec& indices, const arma::Col<T>& data)
+  {
+    for (size_t i = 0; i < indices.n_rows; i++)
+    {
+      element[indices[i] + capacity] = data[i];
+    }
+    // update the total tree with bottom-up technique.
+    for (size_t i = capacity - 1; i > 0; i--)
+    {
+      element[i] = element[2 * i] + element[2 * i + 1];
+    }
+  }
+
+  /**
+   * Get the data array with idx.
+   *
+   * @param idx The array idx to get data.
+   */
+  T Get(size_t idx)
+  {
+    idx += capacity;
+    return element[idx];
+  }
+
+  /**
+   * Help function for the `sum` function
+   *
+   * @param start The starting position of subsequence.
+   * @param end The end position of subsequence.
+   * @param node Reference position.
+   * @param nodeStart Starting position of reference segment.
+   * @param nodeEnd End position of reference segment.
+   */
+  T SumHelper(const size_t start,
+              const size_t end,
+              const size_t node,
+              const size_t nodeStart,
+              const size_t nodeEnd)
+  {
+    if (start == nodeStart && end == nodeEnd)
+    {
+      return element[node];
+    }
+    size_t mid = (nodeStart + nodeEnd) / 2;
+    if (end <= mid)
+    {
+      return SumHelper(start, end, 2 * node, nodeStart, mid);
+    }
+    else
+    {
+      if (mid + 1 <= start)
+      {
+        return SumHelper(start, end, 2 * node + 1, mid + 1 , nodeEnd);
+      }
+      else
+      {
+        return SumHelper(start, mid, 2 * node, nodeStart, mid) +
+            SumHelper(mid + 1, end, 2 * node + 1, mid + 1 , nodeEnd);
+      }
+    }
+  }
+
+  /**
+   * Calculate the sum of contiguous subsequence of the array.
+   *
+   * @param start The starting position of subsequence.
+   * @param end The end position of subsequence.
+   */
+  T Sum(const size_t start, size_t end)
+  {
+    end -= 1;
+    return SumHelper(start, end, 1, 0, capacity - 1);
+  }
+
+  /**
+   * Shortcut for calculating the sum of whole array.
+   */
+  T Sum()
+  {
+    return Sum(0, capacity);
+  }
+
+  /**
+   * Find the highest index `idx` in the array such that
+   * sum(arr[0] + arr[1] + ... + arr[idx]) <= mass.
+   *
+   * @param mass The upper bound of segment array sum.
+   */
+  size_t FindPrefixSum(T mass)
+  {
+    size_t idx = 1;
+    while (idx < capacity)
+    {
+      if (element[2 * idx] > mass)
+      {
+        idx = 2 * idx;
+      }
+      else
+      {
+        mass -= element[2 * idx];
+        idx = 2 * idx + 1;
+      }
+    }
+    return idx - capacity;
+  }
+
+ private:
+  //! The capacity of the data array.
+  size_t capacity;
+
+  //! Double size of capacity, maintain the segment sum of data.
+  std::vector<T> element;
+};
+
+} // namespace rl
+} // namespace mlpack
+
+#endif

src/mlpack/methods/softmax_regression/softmax_regression_main.cpp

@@ -200,10 +200,6 @@ void TestClassifyAcc(size_t numClasses, const Model& model)
   arma::Row<size_t> predictLabels;
   model.Classify(testData, predictLabels);
 
-  // Save predictions, if desired.
-  if (CLI::HasParam("predictions"))
-    CLI::GetParam<arma::Row<size_t>>("predictions") = std::move(predictLabels);
-
   // Calculate accuracy, if desired.
   if (CLI::HasParam("test_labels"))
   {
@@ -242,6 +238,9 @@ void TestClassifyAcc(size_t numClasses, const Model& model)
         << (totalBingo) / static_cast<double>(predictLabels.n_elem) << " ("
         << totalBingo << " of " << predictLabels.n_elem << ")." << endl;
   }
+  // Save predictions, if desired.
+  if (CLI::HasParam("predictions"))
+    CLI::GetParam<arma::Row<size_t>>("predictions") = std::move(predictLabels);
 }
 
 template<typename Model>

src/mlpack/tests/CMakeLists.txt

@@ -99,6 +99,7 @@ add_executable(mlpack_test
   sparse_coding_test.cpp
   spill_tree_test.cpp
   split_data_test.cpp
+  sumtree_test.cpp
   svd_batch_test.cpp
   svd_incremental_test.cpp
   svdplusplus_test.cpp
@@ -121,6 +122,9 @@ add_executable(mlpack_test
   main_tests/det_test.cpp
   main_tests/decision_tree_test.cpp
   main_tests/decision_stump_test.cpp
+  main_tests/gmm_generate_test.cpp
+  main_tests/gmm_probability_test.cpp
+  main_tests/gmm_train_test.cpp
   main_tests/fastmks_test.cpp
   main_tests/kde_test.cpp
   main_tests/linear_regression_test.cpp
@@ -185,7 +189,7 @@ set(parallel_tests
   "SVDIncrementalTest;SVDBatchTest;"
   "LocalCoordinateCodingTest;FeedForwardNetworkTest;SparseAutoencoderTest;"
   "GMMTest;CFTest;ConvolutionalNetworkTest;HMMTest;LARSTest;"
-  "LogisticRegressionTest;"
+  "LogisticRegressionTest;GmmTrainMainTest;"
   "LinearSVMTest")
 
 # Add tests to the testing framework

src/mlpack/tests/ann_layer_test.cpp

@@ -95,6 +95,7 @@ BOOST_AUTO_TEST_CASE(GradientAddLayerTest)
       model->Predictors() = input;
       model->Responses() = target;
       model->Add<IdentityLayer<> >();
+      model->Add<Linear<> >(10, 10);
       model->Add<Add<> >(10);
       model->Add<LogSoftMax<> >();
     }
@@ -400,6 +401,7 @@ BOOST_AUTO_TEST_CASE(GradientLinearLayerTest)
       model->Predictors() = input;
       model->Responses() = target;
       model->Add<IdentityLayer<> >();
+      model->Add<Linear<> >(10, 10);
       model->Add<Linear<> >(10, 2);
       model->Add<LogSoftMax<> >();
     }
@@ -483,6 +485,7 @@ BOOST_AUTO_TEST_CASE(GradientLinearNoBiasLayerTest)
       model->Predictors() = input;
       model->Responses() = target;
       model->Add<IdentityLayer<> >();
+      model->Add<Linear<> >(10, 10);
       model->Add<LinearNoBias<> >(10, 2);
       model->Add<LogSoftMax<> >();
     }
@@ -585,6 +588,7 @@ BOOST_AUTO_TEST_CASE(GradientFlexibleReLULayerTest)
 
       model->Predictors() = input;
       model->Responses() = target;
+      model->Add<Linear<> >(2, 2);
       model->Add<LinearNoBias<> >(2, 5);
       model->Add<FlexibleReLU<> >(0.05);
       model->Add<LogSoftMax<> >();
@@ -1307,6 +1311,7 @@ BOOST_AUTO_TEST_CASE(GradientConcatLayerTest)
       model->Predictors() = input;
       model->Responses() = target;
       model->Add<IdentityLayer<> >();
+      model->Add<Linear<> >(10, 10);
 
       concat = new Concat<>(true);
       concat->Add<Linear<> >(10, 2);
@@ -1578,6 +1583,7 @@ BOOST_AUTO_TEST_CASE(GradientBatchNormTest)
       model->Predictors() = input;
       model->Responses() = target;
       model->Add<IdentityLayer<> >();
+      model->Add<Linear<> >(10, 10);
       model->Add<BatchNorm<> >(10);
       model->Add<Linear<> >(10, 2);
       model->Add<LogSoftMax<> >();
@@ -1751,6 +1757,7 @@ BOOST_AUTO_TEST_CASE(GradientTransposedConvolutionLayerTest)
         model = new FFN<NegativeLogLikelihood<>, RandomInitialization>();
         model->Predictors() = input;
         model->Responses() = target;
+        model->Add<Linear<> >(36, 36);
         model->Add<TransposedConvolution<> >(1, 1, 3, 3, 2, 2, 1, 1, 6, 6);
         model->Add<LogSoftMax<> >();
       }
@@ -1868,6 +1875,7 @@ BOOST_AUTO_TEST_CASE(GradientAtrousConvolutionLayerTest)
       model = new FFN<NegativeLogLikelihood<>, RandomInitialization>();
       model->Predictors() = input;
       model->Responses() = target;
+      model->Add<Linear<> >(36, 36);
       model->Add<AtrousConvolution<> >(1, 1, 3, 3, 1, 1, 0, 0, 6, 6, 2, 2);
       model->Add<LogSoftMax<> >();
     }
@@ -1947,6 +1955,7 @@ BOOST_AUTO_TEST_CASE(GradientLayerNormTest)
       model->Predictors() = input;
       model->Responses() = target;
       model->Add<IdentityLayer<> >();
+      model->Add<Linear<> >(10, 10);
       model->Add<LayerNorm<> >(10);
       model->Add<Linear<> >(10, 2);
       model->Add<LogSoftMax<> >();
@@ -2338,6 +2347,97 @@ BOOST_AUTO_TEST_CASE(SimpleResidualLayerTest)
   delete linearB;
 }
 
+/**
+ * Simple Highway module test.
+ */
+BOOST_AUTO_TEST_CASE(SimpleHighwayLayerTest)
+{
+  arma::mat outputA, outputB, input, deltaA, deltaB;
+  Sequential<>* sequential = new Sequential<>(true);
+  Highway<>* highway = new Highway<>(10, true);
+  highway->Parameters().zeros();
+  highway->Reset();
+
+  Linear<>* linearA = new Linear<>(10, 10);
+  linearA->Parameters().randu();
+  linearA->Reset();
+  Linear<>* linearB = new Linear<>(10, 10);
+  linearB->Parameters().randu();
+  linearB->Reset();
+
+  // Add the same layers (with the same parameters) to both Sequential and
+  // Highway object.
+  highway->Add(linearA);
+  highway->Add(linearB);
+  sequential->Add(linearA);
+  sequential->Add(linearB);
+
+  // Test the Forward function (pass the same input to both).
+  input = arma::randu(10, 1);
+  sequential->Forward(std::move(input), std::move(outputA));
+  highway->Forward(std::move(input), std::move(outputB));
+
+  CheckMatrices(outputB, input * 0.5 + outputA * 0.5);
+
+  delete sequential;
+  delete highway;
+  delete linearA;
+  delete linearB;
+}
+
+/**
+ * Sequential layer numerical gradient test.
+ */
+BOOST_AUTO_TEST_CASE(GradientHighwayLayerTest)
+{
+  // Linear function gradient instantiation.
+  struct GradientFunction
+  {
+    GradientFunction()
+    {
+      input = arma::randu(5, 1);
+      target = arma::mat("1");
+
+      model = new FFN<NegativeLogLikelihood<>, NguyenWidrowInitialization>();
+      model->Predictors() = input;
+      model->Responses() = target;
+      model->Add<IdentityLayer<> >();
+      model->Add<Linear<> >(5, 10);
+
+      highway = new Highway<>(10);
+      highway->Add<Linear<> >(10, 10);
+      highway->Add<ReLULayer<> >();
+      highway->Add<Linear<> >(10, 10);
+      highway->Add<ReLULayer<> >();
+
+      model->Add(highway);
+      model->Add<Linear<> >(10, 2);
+      model->Add<LogSoftMax<> >();
+    }
+
+    ~GradientFunction()
+    {
+      highway->DeleteModules();
+      delete model;
+    }
+
+    double Gradient(arma::mat& gradient) const
+    {
+      double error = model->Evaluate(model->Parameters(), 0, 1);
+      model->Gradient(model->Parameters(), 0, gradient, 1);
+      return error;
+    }
+
+    arma::mat& Parameters() { return model->Parameters(); }
+
+    FFN<NegativeLogLikelihood<>, NguyenWidrowInitialization>* model;
+    Highway<>* highway;
+    arma::mat input, target;
+  } function;
+
+  BOOST_REQUIRE_LE(CheckGradient(function), 1e-4);
+}
+
 /**
  * Sequential layer numerical gradient test.
  */
@@ -2355,7 +2455,7 @@ BOOST_AUTO_TEST_CASE(GradientSequentialLayerTest)
       model->Predictors() = input;
       model->Responses() = target;
       model->Add<IdentityLayer<> >();
-
+      model->Add<Linear<> >(10, 10);
       sequential = new Sequential<>();
       sequential->Add<Linear<> >(10, 10);
       sequential->Add<ReLULayer<> >();

src/mlpack/tests/decision_tree_test.cpp

@@ -288,7 +288,8 @@ BOOST_AUTO_TEST_CASE(BestBinaryNumericSplitSimpleSplitTest)
   // Call the method to do the splitting.
   const double bestGain = GiniGain::Evaluate<false>(labels, 2, weights);
   const double gain = BestBinaryNumericSplit<GiniGain>::SplitIfBetter<false>(
-      bestGain, values, labels, 2, weights, 3, 1e-7, classProbabilities, aux);
+      bestGain, values, labels, 2, weights, 3, 1e-7, classProbabilities,
+      aux);
   const double weightedGain =
       BestBinaryNumericSplit<GiniGain>::SplitIfBetter<true>(bestGain, values,
       labels, 2, weights, 3, 1e-7, classProbabilities, aux);
@@ -325,7 +326,8 @@ BOOST_AUTO_TEST_CASE(BestBinaryNumericSplitMinSamplesTest)
   // Call the method to do the splitting.
   const double bestGain = GiniGain::Evaluate<false>(labels, 2, weights);
   const double gain = BestBinaryNumericSplit<GiniGain>::SplitIfBetter<false>(
-      bestGain, values, labels, 2, weights, 8, 1e-7, classProbabilities, aux);
+      bestGain, values, labels, 2, weights, 8, 1e-7, classProbabilities,
+      aux);
   // This should make no difference because it won't split at all.
   const double weightedGain =
       BestBinaryNumericSplit<GiniGain>::SplitIfBetter<true>(bestGain, values,
@@ -360,7 +362,8 @@ BOOST_AUTO_TEST_CASE(BestBinaryNumericSplitNoGainTest)
   // Call the method to do the splitting.
   const double bestGain = GiniGain::Evaluate<false>(labels, 2, weights);
   const double gain = BestBinaryNumericSplit<GiniGain>::SplitIfBetter<false>(
-      bestGain, values, labels, 2, weights, 10, 1e-7, classProbabilities, aux);
+      bestGain, values, labels, 2, weights, 10, 1e-7, classProbabilities,
+      aux);
 
   // Make sure there was no split.
   BOOST_REQUIRE_EQUAL(gain, DBL_MAX);
@@ -453,8 +456,8 @@ BOOST_AUTO_TEST_CASE(AllCategoricalSplitNoGainTest)
   // Call the method to do the splitting.
   const double bestGain = GiniGain::Evaluate<false>(labels, 3, weights);
   const double gain = AllCategoricalSplit<GiniGain>::SplitIfBetter<false>(
-      bestGain, values, 10, labels, 3, weights, 10, 1e-7, classProbabilities,
-      aux);
+      bestGain, values, 10, labels, 3, weights, 10, 1e-7,
+      classProbabilities, aux);
   const double weightedGain =
       AllCategoricalSplit<GiniGain>::SplitIfBetter<true>(bestGain, values, 10,
       labels, 3, weights, 10, 1e-7, classProbabilities, aux);
@@ -1213,4 +1216,34 @@ BOOST_AUTO_TEST_CASE(DecisionTreeCategoricalTrainReturnEntropy)
   BOOST_REQUIRE_EQUAL(std::isfinite(entropy), true);
 }
 
+/**
+ * Make sure different maximum depth values give different numbers of children.
+ */
+BOOST_AUTO_TEST_CASE(DifferentMaximumDepthTest)
+{
+  arma::mat dataset;
+  arma::Row<size_t> labels;
+  data::Load("vc2.csv", dataset);
+  data::Load("vc2_labels.txt", labels);
+
+  DecisionTree<> d(dataset, labels, 3, 10, 1e-7, 1);
+
+  DecisionTree<> d1(dataset, labels, 3, 10, 1e-7, 2);
+
+  DecisionTree<> d2(dataset, labels, 3, 10, 1e-7);
+
+  // Now require that we have zero children.
+  BOOST_REQUIRE_EQUAL(d.NumChildren(), 0);
+
+  // Now require that we have two children.
+  BOOST_REQUIRE_EQUAL(d1.NumChildren(), 2);
+  BOOST_REQUIRE_EQUAL(d1.Child(0).NumChildren(), 0);
+  BOOST_REQUIRE_EQUAL(d1.Child(1).NumChildren(), 0);
+
+  // Now require that we have two children.
+  BOOST_REQUIRE_EQUAL(d2.NumChildren(), 2);
+  BOOST_REQUIRE_EQUAL(d2.Child(0).NumChildren(), 2);
+  BOOST_REQUIRE_EQUAL(d2.Child(1).NumChildren(), 2);
+}
+
 BOOST_AUTO_TEST_SUITE_END();

src/mlpack/tests/feedforward_network_test.cpp

@@ -29,47 +29,17 @@ using namespace mlpack::ann;
 BOOST_AUTO_TEST_SUITE(FeedForwardNetworkTest);
 
 /**
- * Train and evaluate a vanilla network with the specified structure.
+ * Train and evaluate a model with the specified structure.
  */
-template<typename MatType = arma::mat>
-void BuildVanillaNetwork(MatType& trainData,
-                         MatType& trainLabels,
-                         MatType& testData,
-                         MatType& testLabels,
-                         const size_t outputSize,
-                         const size_t hiddenLayerSize,
-                         const size_t maxEpochs,
-                         const double classificationErrorThreshold)
+template<typename MatType = arma::mat, typename ModelType>
+void TestNetwork(ModelType& model,
+                 MatType& trainData,
+                 MatType& trainLabels,
+                 MatType& testData,
+                 MatType& testLabels,
+                 const size_t maxEpochs,
+                 const double classificationErrorThreshold)
 {
-  /*
-   * Construct a feed forward network with trainData.n_rows input nodes,
-   * hiddenLayerSize hidden nodes and trainLabels.n_rows output nodes. The
-   * network structure looks like:
-   *
-   *  Input         Hidden        Output
-   *  Layer         Layer         Layer
-   * +-----+       +-----+       +-----+
-   * |     |       |     |       |     |
-   * |     +------>|     +------>|     |
-   * |     |     +>|     |     +>|     |
-   * +-----+     | +--+--+     | +-----+
-   *             |             |
-   *  Bias       |  Bias       |
-   *  Layer      |  Layer      |
-   * +-----+     | +-----+     |
-   * |     |     | |     |     |
-   * |     +-----+ |     +-----+
-   * |     |       |     |
-   * +-----+       +-----+
-   */
-
-  FFN<NegativeLogLikelihood<> > model;
-  model.Add<Linear<> >(trainData.n_rows, hiddenLayerSize);
-  model.Add<SigmoidLayer<> >();
-  model.Add<Linear<> >(hiddenLayerSize, outputSize);
-  model.Add<LogSoftMax<> >();
-
-  // RMSProp opt(0.01, 32, 0.88, 1e-8, maxEpochs * trainData.n_cols, -1);
   ens::RMSProp opt(0.01, 32, 0.88, 1e-8, maxEpochs * trainData.n_cols, -1);
   model.Train(trainData, trainLabels, opt);
 
@@ -115,11 +85,38 @@ BOOST_AUTO_TEST_CASE(VanillaNetworkTest)
   arma::mat testLabels = testData.row(testData.n_rows - 1);
   testData.shed_row(testData.n_rows - 1);
 
+  /*
+   * Construct a feed forward network with trainData.n_rows input nodes,
+   * hiddenLayerSize hidden nodes and trainLabels.n_rows output nodes. The
+   * network structure looks like:
+   *
+   *  Input         Hidden        Output
+   *  Layer         Layer         Layer
+   * +-----+       +-----+       +-----+
+   * |     |       |     |       |     |
+   * |     +------>|     +------>|     |
+   * |     |     +>|     |     +>|     |
+   * +-----+     | +--+--+     | +-----+
+   *             |             |
+   *  Bias       |  Bias       |
+   *  Layer      |  Layer      |
+   * +-----+     | +-----+     |
+   * |     |     | |     |     |
+   * |     +-----+ |     +-----+
+   * |     |       |     |
+   * +-----+       +-----+
+   */
+
+  FFN<NegativeLogLikelihood<> > model;
+  model.Add<Linear<> >(trainData.n_rows, 8);
+  model.Add<SigmoidLayer<> >();
+  model.Add<Linear<> >(8, 3);
+  model.Add<LogSoftMax<> >();
+
   // Vanilla neural net with logistic activation function.
-  // Because 92 percent of the patients are not hyperthyroid the neural
+  // Because 92% of the patients are not hyperthyroid the neural
   // network must be significant better than 92%.
-  BuildVanillaNetwork<>
-      (trainData, trainLabels, testData, testLabels, 3, 8, 10, 0.1);
+  TestNetwork<>(model, trainData, trainLabels, testData, testLabels, 10, 0.1);
 
   arma::mat dataset;
   dataset.load("mnist_first250_training_4s_and_9s.arm");
@@ -132,9 +129,13 @@ BOOST_AUTO_TEST_CASE(VanillaNetworkTest)
   labels.submat(0, labels.n_cols / 2, 0, labels.n_cols - 1).fill(1);
   labels += 1;
 
+  FFN<NegativeLogLikelihood<> > model1;
+  model1.Add<Linear<> >(dataset.n_rows, 10);
+  model1.Add<SigmoidLayer<> >();
+  model1.Add<Linear<> >(10, 2);
+  model1.Add<LogSoftMax<> >();
   // Vanilla neural net with logistic activation function.
-  BuildVanillaNetwork<>
-      (dataset, labels, dataset, labels, 2, 10, 10, 0.2);
+  TestNetwork<>(model1, dataset, labels, dataset, labels, 10, 0.2);
 }
 
 BOOST_AUTO_TEST_CASE(ForwardBackwardTest)
@@ -214,18 +215,23 @@ BOOST_AUTO_TEST_CASE(ForwardBackwardTest)
 }
 
 /**
- * Train and evaluate a Dropout network with the specified structure.
+ * Train the dropout network on a larger dataset.
  */
-template<typename MatType = arma::mat>
-void BuildDropoutNetwork(MatType& trainData,
-                         MatType& trainLabels,
-                         MatType& testData,
-                         MatType& testLabels,
-                         const size_t outputSize,
-                         const size_t hiddenLayerSize,
-                         const size_t maxEpochs,
-                         const double classificationErrorThreshold)
+BOOST_AUTO_TEST_CASE(DropoutNetworkTest)
 {
+  // Load the dataset.
+  arma::mat trainData;
+  data::Load("thyroid_train.csv", trainData, true);
+
+  arma::mat trainLabels = trainData.row(trainData.n_rows - 1);
+  trainData.shed_row(trainData.n_rows - 1);
+
+  arma::mat testData;
+  data::Load("thyroid_test.csv", testData, true);
+
+  arma::mat testLabels = testData.row(testData.n_rows - 1);
+  testData.shed_row(testData.n_rows - 1);
+
   /*
    * Construct a feed forward network with trainData.n_rows input nodes,
    * hiddenLayerSize hidden nodes and trainLabels.n_rows output nodes. The
@@ -249,44 +255,70 @@ void BuildDropoutNetwork(MatType& trainData,
    */
 
   FFN<NegativeLogLikelihood<> > model;
-  model.Add<Linear<> >(trainData.n_rows, hiddenLayerSize);
+  model.Add<Linear<> >(trainData.n_rows, 8);
   model.Add<SigmoidLayer<> >();
   model.Add<Dropout<> >();
-  model.Add<Linear<> >(hiddenLayerSize, outputSize);
+  model.Add<Linear<> >(8, 3);
   model.Add<LogSoftMax<> >();
 
-  ens::RMSProp opt(0.01, 32, 0.88, 1e-8, maxEpochs * trainData.n_cols, -1);
+  // Vanilla neural net with logistic activation function.
+  // Because 92% of the patients are not hyperthyroid the neural
+  // network must be significant better than 92%.
+  TestNetwork<>(model, trainData, trainLabels, testData, testLabels, 10, 0.1);
+  arma::mat dataset;
+  dataset.load("mnist_first250_training_4s_and_9s.arm");
 
-  model.Train(trainData, trainLabels, opt);
-
-  MatType predictionTemp;
-  model.Predict(testData, predictionTemp);
-  MatType prediction = arma::zeros<MatType>(1, predictionTemp.n_cols);
-
-  for (size_t i = 0; i < predictionTemp.n_cols; ++i)
+  // Normalize each point since these are images.
+  for (size_t i = 0; i < dataset.n_cols; ++i)
   {
-    prediction(i) = arma::as_scalar(arma::find(
-        arma::max(predictionTemp.col(i)) == predictionTemp.col(i), 1)) + 1;
+    dataset.col(i) /= norm(dataset.col(i), 2);
   }
 
-  size_t error = 0;
-  for (size_t i = 0; i < testData.n_cols; i++)
-  {
-    if (int(arma::as_scalar(prediction.col(i))) ==
-        int(arma::as_scalar(testLabels.col(i))))
-    {
-      error++;
-    }
-  }
+  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;
 
-  double classificationError = 1 - double(error) / testData.n_cols;
-  BOOST_REQUIRE_LE(classificationError, classificationErrorThreshold);
+  FFN<NegativeLogLikelihood<> > model1;
+  model1.Add<Linear<> >(dataset.n_rows, 10);
+  model1.Add<SigmoidLayer<> >();
+  model.Add<Dropout<> >();
+  model1.Add<Linear<> >(10, 2);
+  model1.Add<LogSoftMax<> >();
+  // Vanilla neural net with logistic activation function.
+  TestNetwork<>(model1, dataset, labels, dataset, labels, 10, 0.2);
 }
 
 /**
- * Train the dropout network on a larger dataset.
+ * Train the highway network on a larger dataset.
  */
-BOOST_AUTO_TEST_CASE(DropoutNetworkTest)
+BOOST_AUTO_TEST_CASE(HighwayNetworkTest)
+{
+  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, 10);
+  Highway<>* highway = new Highway<>(10, true);
+  highway->Add<Linear<> >(10, 10);
+  highway->Add<SigmoidLayer<> >();
+  model.Add(highway);
+  model.Add<Linear<> >(10, 2);
+  model.Add<LogSoftMax<> >();
+  TestNetwork<>(model, dataset, labels, dataset, labels, 10, 0.2);
+}
+
+/**
+ * Train the DropConnect network on a larger dataset.
+ */
+BOOST_AUTO_TEST_CASE(DropConnectNetworkTest)
 {
   // Load the dataset.
   arma::mat trainData;
@@ -301,42 +333,6 @@ BOOST_AUTO_TEST_CASE(DropoutNetworkTest)
   arma::mat testLabels = testData.row(testData.n_rows - 1);
   testData.shed_row(testData.n_rows - 1);
 
-  // Vanilla neural net with logistic activation function.
-  // Because 92 percent of the patients are not hyperthyroid the neural
-  // network must be significant better than 92%.
-  BuildDropoutNetwork<>
-      (trainData, trainLabels, testData, testLabels, 3, 8, 10, 0.1);
-
-  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;
-
-  // Vanilla neural net with logistic activation function.
-  BuildDropoutNetwork<>
-      (dataset, labels, dataset, labels, 2, 10, 10, 0.2);
-}
-
-/**
- * Train and evaluate a DropConnect network(with a baselayer) with the
- * specified structure.
- */
-template<typename MatType = arma::mat>
-void BuildDropConnectNetwork(MatType& trainData,
-                             MatType& trainLabels,
-                             MatType& testData,
-                             MatType& testLabels,
-                             const size_t outputSize,
-                             const size_t hiddenLayerSize,
-                             const size_t maxEpochs,
-                             const double classificationErrorThreshold)
-{
  /*
   *  Construct a feed forward network with trainData.n_rows input nodes,
   *  hiddenLayerSize hidden nodes and trainLabels.n_rows output nodes. The
@@ -362,62 +358,15 @@ void BuildDropConnectNetwork(MatType& trainData,
   */
 
   FFN<NegativeLogLikelihood<> > model;
-  model.Add<Linear<> >(trainData.n_rows, hiddenLayerSize);
+  model.Add<Linear<> >(trainData.n_rows, 8);
   model.Add<SigmoidLayer<> >();
-  model.Add<DropConnect<> >(hiddenLayerSize, outputSize);
+  model.Add<DropConnect<> >(8, 3);
   model.Add<LogSoftMax<> >();
 
-  ens::RMSProp opt(0.01, 32, 0.88, 1e-8, maxEpochs * trainData.n_cols, -1);
-
-  model.Train(trainData, trainLabels, opt);
-
-  MatType predictionTemp;
-  model.Predict(testData, predictionTemp);
-  MatType prediction = arma::zeros<MatType>(1, predictionTemp.n_cols);
-
-  for (size_t i = 0; i < predictionTemp.n_cols; ++i)
-  {
-    prediction(i) = arma::as_scalar(arma::find(
-        arma::max(predictionTemp.col(i)) == predictionTemp.col(i), 1)) + 1;
-  }
-
-  size_t error = 0;
-  for (size_t i = 0; i < testData.n_cols; i++)
-  {
-    if (int(arma::as_scalar(prediction.col(i))) ==
-        int(arma::as_scalar(testLabels.col(i))))
-    {
-      error++;
-    }
-  }
-
-  double classificationError = 1 - double(error) / testData.n_cols;
-  BOOST_REQUIRE_LE(classificationError, classificationErrorThreshold);
-}
-
-/**
- * Train the dropconnect network on a larger dataset.
- */
-BOOST_AUTO_TEST_CASE(DropConnectNetworkTest)
-{
-  // Load the dataset.
-  arma::mat trainData;
-  data::Load("thyroid_train.csv", trainData, true);
-
-  arma::mat trainLabels = trainData.row(trainData.n_rows - 1);
-  trainData.shed_row(trainData.n_rows - 1);
-
-  arma::mat testData;
-  data::Load("thyroid_test.csv", testData, true);
-
-  arma::mat testLabels = testData.row(testData.n_rows - 1);
-  testData.shed_row(testData.n_rows - 1);
-
   // Vanilla neural net with logistic activation function.
-  // Because 92 percent of the patients are not hyperthyroid the neural
+  // Because 92% of the patients are not hyperthyroid the neural
   // network must be significant better than 92%.
-  BuildDropConnectNetwork<>
-      (trainData, trainLabels, testData, testLabels, 3, 8, 10, 0.1);
+  TestNetwork<>(model, trainData, trainLabels, testData, testLabels, 10, 0.1);
 
   arma::mat dataset;
   dataset.load("mnist_first250_training_4s_and_9s.arm");
@@ -430,9 +379,13 @@ BOOST_AUTO_TEST_CASE(DropConnectNetworkTest)
   labels.submat(0, labels.n_cols / 2, 0, labels.n_cols - 1).fill(1);
   labels += 1;
 
+  FFN<NegativeLogLikelihood<> > model1;
+  model1.Add<Linear<> >(dataset.n_rows, 10);
+  model1.Add<SigmoidLayer<> >();
+  model1.Add<DropConnect<> >(10, 2);
+  model1.Add<LogSoftMax<> >();
   // Vanilla neural net with logistic activation function.
-  BuildDropConnectNetwork<>
-      (dataset, labels, dataset, labels, 2, 10, 10, 0.2);
+  TestNetwork<>(model1, dataset, labels, dataset, labels, 10, 0.2);
 }
 
 /**
@@ -470,7 +423,7 @@ BOOST_AUTO_TEST_CASE(SerializationTest)
   testData.shed_row(testData.n_rows - 1);
 
   // Vanilla neural net with logistic activation function.
-  // Because 92 percent of the patients are not hyperthyroid the neural
+  // Because 92% of the patients are not hyperthyroid the neural
   // network must be significant better than 92%.
   FFN<NegativeLogLikelihood<> > model;
   model.Add<Linear<> >(trainData.n_rows, 8);
@@ -598,7 +551,7 @@ BOOST_AUTO_TEST_CASE(FFNTrainReturnObjective)
   testData.shed_row(testData.n_rows - 1);
 
   // Vanilla neural net with logistic activation function.
-  // Because 92 percent of the patients are not hyperthyroid the neural
+  // Because 92% of the patients are not hyperthyroid the neural
   // network must be significantly better than 92%.
   FFN<NegativeLogLikelihood<> > model;
   model.Add<Linear<> >(trainData.n_rows, 8);

src/mlpack/tests/gan_test.cpp

@@ -246,4 +246,63 @@ BOOST_AUTO_TEST_CASE(GANMNISTTest)
   Log::Info << "Output generated!" << std::endl;
 }
 
+/*
+ * Create GAN network and test for memory sharing
+ * between discriminator and gan predictors.
+ */
+BOOST_AUTO_TEST_CASE(GANMemorySharingTest)
+{
+  size_t generatorHiddenLayerSize = 8;
+  size_t discriminatorHiddenLayerSize = 8;
+  size_t generatorOutputSize = 1;
+  size_t discriminatorOutputSize = 1;
+  size_t discriminatorPreTrain = 0;
+  size_t batchSize = 8;
+  size_t noiseDim = 1;
+  size_t generatorUpdateStep = 1;
+  double multiplier = 1;
+
+  arma::mat trainData(1, 10000);
+  trainData.imbue( [&]() { return arma::as_scalar(RandNormal(4, 0.5));});
+  trainData = arma::sort(trainData);
+
+  // Create the Discriminator network
+  FFN<SigmoidCrossEntropyError<> > discriminator;
+  discriminator.Add<Linear<> > (
+      generatorOutputSize, discriminatorHiddenLayerSize * 2);
+  discriminator.Add<ReLULayer<> >();
+  discriminator.Add<Linear<> > (
+      discriminatorHiddenLayerSize * 2, discriminatorHiddenLayerSize * 2);
+  discriminator.Add<ReLULayer<> >();
+  discriminator.Add<Linear<> > (
+      discriminatorHiddenLayerSize * 2, discriminatorHiddenLayerSize * 2);
+  discriminator.Add<ReLULayer<> >();
+  discriminator.Add<Linear<> > (
+      discriminatorHiddenLayerSize * 2, discriminatorOutputSize);
+
+  // Create the Generator network
+  FFN<SigmoidCrossEntropyError<> > generator;
+  generator.Add<Linear<> >(noiseDim, generatorHiddenLayerSize);
+  generator.Add<SoftPlusLayer<> >();
+  generator.Add<Linear<> >(generatorHiddenLayerSize, generatorOutputSize);
+
+  // Create GAN
+  GaussianInitialization gaussian(0, 0.1);
+  std::function<double ()> noiseFunction = [](){ return math::Random(-8, 8) +
+      math::RandNormal(0, 1) * 0.01;};
+  GAN<FFN<SigmoidCrossEntropyError<> >,
+      GaussianInitialization,
+      std::function<double()> >
+  gan(trainData, generator, discriminator, gaussian, noiseFunction,
+      noiseDim, batchSize, generatorUpdateStep, discriminatorPreTrain,
+      multiplier);
+
+  CheckMatrices(gan.Predictors().head_cols(trainData.n_cols), trainData);
+  CheckMatrices(gan.Predictors(), gan.Discriminator().Predictors());
+  gan.Shuffle();
+  CheckMatrices(gan.Predictors(), gan.Discriminator().Predictors());
+  CheckMatricesNotEqual(gan.Predictors().head_cols(trainData.n_cols),
+      trainData);
+}
+
 BOOST_AUTO_TEST_SUITE_END();

src/mlpack/tests/main_tests/decision_tree_test.cpp

@@ -41,6 +41,12 @@ struct DecisionTreeTestFixture
   }
 };
 
+void ResetDTSettings()
+{
+  CLI::ClearSettings();
+  CLI::RestoreSettings(testName);
+}
+
 BOOST_FIXTURE_TEST_SUITE(DecisionTreeMainTest,
                          DecisionTreeTestFixture);
 
@@ -167,6 +173,35 @@ BOOST_AUTO_TEST_CASE(DecisionTreeMinimumLeafSizeTest)
   Log::Fatal.ignoreInput = false;
 }
 
+/**
+ * Make sure maximum depth is always a non-negative number.
+ */
+BOOST_AUTO_TEST_CASE(DecisionTreeNonNegativeMaximumDepthTest)
+{
+  arma::mat inputData;
+  DatasetInfo info;
+  if (!data::Load("braziltourism.arff", inputData, info))
+    BOOST_FAIL("Cannot load train dataset braziltourism.arff!");
+
+  arma::Row<size_t> labels;
+  if (!data::Load("braziltourism_labels.txt", labels))
+    BOOST_FAIL("Cannot load labels for braziltourism_labels.txt");
+
+  // Initialize an all-ones weight matrix.
+  arma::mat weights(1, labels.n_cols, arma::fill::ones);
+
+  // Input training data.
+  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("labels", std::move(labels));
+  SetInputParam("weights", std::move(weights));
+
+  SetInputParam("maximum_depth", (int) -1); // Invalid.
+
+  Log::Fatal.ignoreInput = true;
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+  Log::Fatal.ignoreInput = false;
+}
+
 /**
  * Make sure minimum gain split is always a fraction in range [0,1].
  */
@@ -419,4 +454,55 @@ BOOST_AUTO_TEST_CASE(DecisionModelCategoricalReuseTest)
   CheckMatrices(probabilities, CLI::GetParam<arma::mat>("probabilities"));
 }
 
+/**
+ * Check that different maximum depths give different results.
+ */
+BOOST_AUTO_TEST_CASE(DecisionTreeMaximumDepthTest)
+{
+  arma::mat inputData;
+  DatasetInfo info;
+  if (!data::Load("vc2.csv", inputData, info))
+    BOOST_FAIL("Cannot load train dataset vc2.csv!");
+
+  arma::Row<size_t> labels;
+  if (!data::Load("vc2_labels.txt", labels))
+    BOOST_FAIL("Cannot load labels for vc2_labels.txt");
+
+  // Initialize an all-ones weight matrix.
+  arma::mat weights(1, labels.n_cols, arma::fill::ones);
+
+  arma::mat testData;
+  if (!data::Load("vc2_test.csv", testData, info))
+    BOOST_FAIL("Cannot load test dataset vc2.csv!");
+
+  // Input training data.
+  SetInputParam("training", std::make_tuple(info, inputData));
+  SetInputParam("labels", labels);
+  SetInputParam("weights", weights);
+  SetInputParam("maximum_depth", (int) 0);
+
+  // Input test data.
+  SetInputParam("test", std::make_tuple(info, testData));
+
+  mlpackMain();
+
+  // Check that number of output points are equal to number of input points.
+  arma::Row<size_t> predictions;
+  predictions = CLI::GetParam<arma::Row<size_t>>("predictions");
+
+  ResetDTSettings();
+
+  // Input training data.
+  SetInputParam("training", std::make_tuple(info, inputData));
+  SetInputParam("labels", std::move(labels));
+  SetInputParam("weights", std::move(weights));
+  SetInputParam("maximum_depth", (int) 4);
+
+  // Input test data.
+  SetInputParam("test", std::make_tuple(info, testData));
+
+  CheckMatricesNotEqual(predictions,
+                        CLI::GetParam<arma::Row<size_t>>("predictions"));
+}
+
 BOOST_AUTO_TEST_SUITE_END();

src/mlpack/tests/main_tests/gmm_generate_test.cpp

@@ -0,0 +1,78 @@
+/**
+ * @file gmm_generate_test.cpp
+ * @author Yashwant Singh
+ *
+ * Test mlpackMain() of gmm_generate_main.cpp.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#define BINDING_TYPE BINDING_TYPE_TEST
+static const std::string testName = "GmmGenerate";
+
+#include <mlpack/core.hpp>
+#include <mlpack/core/util/mlpack_main.hpp>
+#include <mlpack/methods/gmm/gmm_generate_main.cpp>
+
+#include "test_helper.hpp"
+#include <boost/test/unit_test.hpp>
+#include "../test_tools.hpp"
+
+using namespace mlpack;
+
+struct GmmGenerateTestFixture
+{
+ public:
+  GmmGenerateTestFixture()
+  {
+    // Cache in the options for this program.
+    CLI::RestoreSettings(testName);
+  }
+
+  ~GmmGenerateTestFixture()
+  {
+    // Clear the settings.
+    CLI::ClearSettings();
+  }
+};
+
+BOOST_FIXTURE_TEST_SUITE(GmmGenerateMainTest, GmmGenerateTestFixture);
+
+// Checking that Samples must greater than 0.
+BOOST_AUTO_TEST_CASE(GmmGenerateSamplesTest)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  GMM gmm(1, 5);
+  gmm.Train(inputData, 5);
+
+  SetInputParam("input_model", &gmm);
+
+  Log::Fatal.ignoreInput = true;
+  SetInputParam("samples", 0); // Invalid
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+  Log::Fatal.ignoreInput = false;
+}
+
+// Checking dimensionality of output.
+BOOST_AUTO_TEST_CASE(GmmGenerateDimensionality)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  GMM gmm(1, 5);
+  gmm.Train(inputData, 5);
+  SetInputParam("input_model", &gmm);
+  SetInputParam("samples", (int) 10);
+
+  mlpackMain();
+
+  arma::mat output = std::move(CLI::GetParam<arma::mat>("output"));
+
+  BOOST_REQUIRE_EQUAL(output.n_rows, gmm.Dimensionality());
+  BOOST_REQUIRE_EQUAL(output.n_cols, (int) 10);
+}
+
+BOOST_AUTO_TEST_SUITE_END();
+

src/mlpack/tests/main_tests/gmm_probability_test.cpp

@@ -0,0 +1,72 @@
+/**
+ * @file gmm_probability_test.cpp
+ * @author Yashwant Singh
+ *
+ * Test mlpackMain() of gmm_probability_main.cpp.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+
+#define BINDING_TYPE BINDING_TYPE_TEST
+
+static const std::string testName = "GmmProbability";
+
+#include <mlpack/core.hpp>
+#include <mlpack/core/util/mlpack_main.hpp>
+#include <mlpack/methods/gmm/gmm_probability_main.cpp>
+
+#include "test_helper.hpp"
+
+#include <boost/test/unit_test.hpp>
+
+
+using namespace mlpack;
+
+struct GmmProbabilityTestFixture
+{
+ public:
+  GmmProbabilityTestFixture()
+  {
+    // Cache in the options for this program.
+    CLI::RestoreSettings(testName);
+  }
+
+  ~GmmProbabilityTestFixture()
+  {
+    // Clear the settings.
+    CLI::ClearSettings();
+  }
+};
+
+void ResetGmmProbabilitySetting()
+{
+  CLI::ClearSettings();
+  CLI::RestoreSettings(testName);
+}
+
+BOOST_FIXTURE_TEST_SUITE(GmmProbabilityMainTest, GmmProbabilityTestFixture);
+
+// Checking the input and output dimensionality.
+BOOST_AUTO_TEST_CASE(GmmProbabilityDimensionality)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  GMM gmm(1, 5);
+  gmm.Train(std::move(inputData), 5);
+
+  arma::mat inputPoints(5, 5, arma::fill::randu);
+
+  SetInputParam("input", std::move(inputPoints));
+  SetInputParam("input_model", &gmm);
+
+  mlpackMain();
+
+  BOOST_REQUIRE_EQUAL(CLI::GetParam<arma::mat>("output").n_cols, 5);
+  BOOST_REQUIRE_EQUAL(CLI::GetParam<arma::mat>("output").n_rows, 1);
+}
+
+BOOST_AUTO_TEST_SUITE_END();
+

src/mlpack/tests/main_tests/gmm_train_test.cpp

@@ -0,0 +1,490 @@
+/**
+ * @file gmm_train_test.cpp
+ * @author Yashwant Singh
+ *
+ * Test mlpackMain() of gmm_train_main.cpp.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+
+#include<string>
+
+#define BINDING_TYPE BINDING_TYPE_TEST
+static const std::string testName = "GmmTrain";
+
+#include <mlpack/core.hpp>
+#include <mlpack/core/util/mlpack_main.hpp>
+#include "test_helper.hpp"
+#include <mlpack/methods/gmm/gmm_train_main.cpp>
+
+#include <boost/test/unit_test.hpp>
+#include "../test_tools.hpp"
+
+using namespace mlpack;
+
+struct GmmTrainTestFixture
+{
+ public:
+  GmmTrainTestFixture()
+  {
+    // Cache in the options for this program.
+    CLI::RestoreSettings(testName);
+  }
+
+  ~GmmTrainTestFixture()
+  {
+    // Clear the settings.
+    bindings::tests::CleanMemory();
+    CLI::ClearSettings();
+  }
+};
+
+void ResetGmmTrainSetting()
+{
+  CLI::ClearSettings();
+  CLI::RestoreSettings(testName);
+}
+
+
+BOOST_FIXTURE_TEST_SUITE(GmmTrainMainTest, GmmTrainTestFixture);
+
+// To check if the gaussian is positive or not.
+BOOST_AUTO_TEST_CASE(GmmTrainValidGaussianTest)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", 0); // Invalid
+
+  Log::Fatal.ignoreInput = true;
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+  Log::Fatal.ignoreInput = false;
+}
+
+/**
+  * To check if the number of gaussians in the output model is same as
+  * that of input gaussian parameter or not.
+ **/
+BOOST_AUTO_TEST_CASE(GmmTrainOutputModelGaussianTest)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("trials", (int) 2);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+  BOOST_REQUIRE_EQUAL(gmm->Gaussians(), (int) 2);
+}
+
+// Max iterations must be positive.
+BOOST_AUTO_TEST_CASE(GmmTrainMaxIterationsTest)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("trials", (int) 1);
+  SetInputParam("max_iterations", (int)-1); // Invalid.
+
+  Log::Fatal.ignoreInput = true;
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+  Log::Fatal.ignoreInput = false;
+}
+
+// Ensure that Trials must be greater than 0.
+BOOST_AUTO_TEST_CASE(GmmTrainPositiveTrialsTest)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("trials", (int) 0); // Invalid.
+
+  Log::Fatal.ignoreInput = true;
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+  Log::Fatal.ignoreInput = false;
+}
+
+// Checking that percentage is between 0 and 1.
+BOOST_AUTO_TEST_CASE(RefinedStartPercentageTest)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("refined_start", true);
+
+  Log::Fatal.ignoreInput = true;
+  SetInputParam("percentage", (double) 2.0); // Invalid
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+
+  SetInputParam("percentage", (double) -1.0); // Invalid
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+
+  Log::Fatal.ignoreInput = false;
+}
+
+// Samplings must be positive.
+BOOST_AUTO_TEST_CASE(GmmTrainSamplings)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("refined_start", true);
+  SetInputParam("samplings", (int) 0); // Invalid
+
+  Log::Fatal.ignoreInput = true;
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+  Log::Fatal.ignoreInput = false;
+}
+
+// Number of gaussians in the model trained from input model.
+BOOST_AUTO_TEST_CASE(GmmTrainNumberOfGaussian)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", inputData);
+  SetInputParam("gaussians", (int) 2);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  SetInputParam("input_model", gmm);
+
+  CLI::GetSingleton().Parameters()["input"].wasPassed = false;
+
+  SetInputParam("input", std::move(inputData));
+
+  mlpackMain();
+
+  GMM* gmm1 = CLI::GetParam<GMM*>("output_model");
+
+  BOOST_REQUIRE_EQUAL(gmm1->Gaussians(), (int) 2);
+}
+
+// Making sure that enabling no_force_positive doesn't crash.
+BOOST_AUTO_TEST_CASE(GmmTrainNoForcePositiveTest)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", inputData);
+  SetInputParam("gaussians", (int) 1);
+  SetInputParam("no_force_positive", true);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  SetInputParam("input_model", gmm);
+
+  CLI::GetSingleton().Parameters()["input"].wasPassed = false;
+
+  SetInputParam("input", std::move(inputData));
+
+  mlpackMain();
+
+  GMM* gmm1 = CLI::GetParam<GMM*>("output_model");
+
+  BOOST_REQUIRE_EQUAL(gmm1->Gaussians(), (int) 1);
+}
+
+// Ensure that Noise affects the final result.
+BOOST_AUTO_TEST_CASE(GmmTrainNoiseTest)
+{
+  arma::mat inputData;
+  if (!data::Load("data_3d_mixed.txt", inputData))
+    BOOST_FAIL("Unable to load train dataset data_3d_mixed.txt!");
+
+  SetInputParam("input", inputData);
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("noise", (double) 0.0);
+
+  size_t seed = std::time(NULL);
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  ResetGmmTrainSetting();
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("noise", (double) 100.0);
+
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm1 = CLI::GetParam<GMM*>("output_model");
+
+  arma::uvec sortedIndices = sort_index(gmm->Weights());
+
+  for (size_t k = 0; k < sortedIndices.n_elem; k++)
+  {
+    BOOST_REQUIRE(arma::norm(gmm->Component(sortedIndices[k]).Mean() -
+                      gmm1->Component(sortedIndices[k]).Mean()) > 1e-50 ||
+                  arma::norm(gmm->Component(sortedIndices[k]).Covariance() -
+                      gmm1->Component(sortedIndices[k]).Covariance()) > 1e-50);
+  }
+}
+
+// Ensure that Trials affects the final result.
+BOOST_AUTO_TEST_CASE(GmmTrainTrialsTest)
+{
+  arma::mat inputData(5, 250, arma::fill::randu);
+
+  SetInputParam("input", inputData);
+  SetInputParam("gaussians", (int) 3);
+  SetInputParam("trials", (int) 1);
+  SetInputParam("max_iterations", (int) 500);
+
+  size_t seed = std::time(NULL);
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  ResetGmmTrainSetting();
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 3);
+  SetInputParam("max_iterations", (int) 500);
+  SetInputParam("trials", (int) 500);
+
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm1 = CLI::GetParam<GMM*>("output_model");
+
+  arma::uvec sortedIndices = sort_index(gmm->Weights());
+
+  for (size_t k = 0; k < sortedIndices.n_elem; k++)
+  {
+    BOOST_REQUIRE(arma::norm(gmm->Component(sortedIndices[k]).Mean() -
+                      gmm1->Component(sortedIndices[k]).Mean()) > 1e-50 ||
+                  arma::norm(gmm->Component(sortedIndices[k]).Covariance() -
+                      gmm1->Component(sortedIndices[k]).Covariance()) > 1e-50);
+  }
+}
+
+// Ensure that Percentage affects the final result when refined_start is true.
+BOOST_AUTO_TEST_CASE(GmmTrainPercentageTest)
+{
+  arma::mat inputData;
+  if (!data::Load("data_3d_mixed.txt", inputData))
+    BOOST_FAIL("Unable to load train dataset data_3d_mixed.txt!");
+
+  SetInputParam("input", inputData);
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("refined_start", true);
+  SetInputParam("percentage", (double) 0.01);
+  SetInputParam("samplings", (int) 1000);
+
+  size_t seed = std::time(NULL);
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  ResetGmmTrainSetting();
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("refined_start", true);
+  SetInputParam("percentage", (double) 0.45);
+  SetInputParam("samplings", (int) 1000);
+
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm1 = CLI::GetParam<GMM*>("output_model");
+
+  arma::uvec sortedIndices = sort_index(gmm->Weights());
+
+  for (size_t k = 0; k < sortedIndices.n_elem; k++)
+  {
+    BOOST_REQUIRE(arma::norm(gmm->Component(sortedIndices[k]).Mean() -
+                      gmm1->Component(sortedIndices[k]).Mean()) > 1e-50 ||
+                  arma::norm(gmm->Component(sortedIndices[k]).Covariance() -
+                      gmm1->Component(sortedIndices[k]).Covariance()) > 1e-50);
+  }
+}
+
+// Ensure that Sampling affects the final result when refined_start is true.
+BOOST_AUTO_TEST_CASE(GmmTrainSamplingsTest)
+{
+  arma::mat inputData;
+  if (!data::Load("data_3d_mixed.txt", inputData))
+    BOOST_FAIL("Unable to load train dataset data_3d_mixed.txt!");
+
+  SetInputParam("input", inputData);
+  SetInputParam("gaussians", (int) 8);
+  SetInputParam("refined_start", true);
+  SetInputParam("trials", (int) 2);
+  SetInputParam("samplings", (int) 10);
+
+  size_t seed = std::time(NULL);
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  ResetGmmTrainSetting();
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 8);
+  SetInputParam("refined_start", true);
+  SetInputParam("trials", (int) 2);
+  SetInputParam("samplings", (int) 5000);
+
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm1 = CLI::GetParam<GMM*>("output_model");
+
+  arma::uvec sortedIndices = sort_index(gmm->Weights());
+
+  for (size_t k = 0; k < sortedIndices.n_elem; k++)
+  {
+    BOOST_REQUIRE(arma::norm(gmm->Component(sortedIndices[k]).Mean() -
+                      gmm1->Component(sortedIndices[k]).Mean()) > 1e-50 ||
+                  arma::norm(gmm->Component(sortedIndices[k]).Covariance() -
+                      gmm1->Component(sortedIndices[k]).Covariance()) > 1e-50);
+  }
+}
+
+// Ensure that tolerance affects the final result.
+BOOST_AUTO_TEST_CASE(GmmTrainToleranceTest)
+{
+  arma::mat inputData;
+  if (!data::Load("data_3d_mixed.txt", inputData))
+    BOOST_FAIL("Unable to load train dataset data_3d_mixed.txt!");
+
+  SetInputParam("input", inputData);
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("tolerance", (double) 1e-8);
+
+  size_t seed = std::time(NULL);
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  ResetGmmTrainSetting();
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("tolerance", (double) 10);
+
+  mlpack::math::randGen.seed((uint32_t) seed);
+  srand((unsigned int) seed);
+  arma::arma_rng::set_seed(seed);
+
+  mlpackMain();
+
+  GMM* gmm1 = CLI::GetParam<GMM*>("output_model");
+
+  arma::uvec sortedIndices = sort_index(gmm->Weights());
+
+  for (size_t k = 0; k < sortedIndices.n_elem; k++)
+  {
+    BOOST_REQUIRE(arma::norm(gmm->Component(sortedIndices[k]).Mean() -
+                      gmm1->Component(sortedIndices[k]).Mean()) > 1e-50 ||
+                  arma::norm(gmm->Component(sortedIndices[k]).Covariance() -
+                      gmm1->Component(sortedIndices[k]).Covariance()) > 1e-50);
+  }
+}
+
+// Ensure that saved model can be used again.
+BOOST_AUTO_TEST_CASE(GmmTrainModelReuseTest)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", inputData);
+  SetInputParam("gaussians", (int) 2);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  SetInputParam("input_model", gmm);
+
+  CLI::GetSingleton().Parameters()["input"].wasPassed = false;
+
+  SetInputParam("input", inputData);
+
+  mlpackMain();
+
+  GMM* gmm1 = CLI::GetParam<GMM*>("output_model");
+
+  SetInputParam("input_model", gmm1);
+
+  CLI::GetSingleton().Parameters()["input"].wasPassed = false;
+
+  SetInputParam("input", std::move(inputData));
+
+  mlpackMain();
+
+  GMM* gmm2 = CLI::GetParam<GMM*>("output_model");
+
+  BOOST_REQUIRE_EQUAL(gmm1, gmm2);
+}
+
+// Ensure that Gmm's covariances are diagonal when diagonal_covariance is true.
+BOOST_AUTO_TEST_CASE(GmmTrainDiagCovariance)
+{
+  arma::mat inputData(5, 10, arma::fill::randu);
+
+  SetInputParam("input", std::move(inputData));
+  SetInputParam("gaussians", (int) 2);
+  SetInputParam("diagonal_covariance", true);
+
+  mlpackMain();
+
+  GMM* gmm = CLI::GetParam<GMM*>("output_model");
+
+  arma::uvec sortedIndices = sort_index(gmm->Weights());
+
+  for (size_t k = 0; k < sortedIndices.n_elem; k++)
+  {
+    arma::mat diagCov(gmm->Component(sortedIndices[k]).Covariance());
+      for (size_t i = 0; i < diagCov.n_rows; i++)
+        for (size_t j = 0; j < diagCov.n_cols; j++)
+          if (i != j && diagCov(i, j) != (double) 0)
+            BOOST_FAIL("Covariance Are Not Diagonal");
+  }
+}
+
+BOOST_AUTO_TEST_SUITE_END();

src/mlpack/tests/main_tests/hoeffding_tree_test.cpp

@@ -66,11 +66,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingTreeOutputDimensionTest)
   size_t testSize = testData.n_cols;
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   mlpackMain();
 
@@ -109,11 +109,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingTreeCategoricalOutputDimensionTest)
   size_t testSize = testData.n_cols;
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   mlpackMain();
 
@@ -157,10 +157,10 @@ BOOST_AUTO_TEST_CASE(HoeffdingTreeLabelLessTest)
   size_t testSize = testData.n_cols;
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   mlpackMain();
 
@@ -191,8 +191,8 @@ BOOST_AUTO_TEST_CASE(HoeffdingTreeLabelLessTest)
   inputData.shed_row(inputData.n_rows - 1);
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
+  SetInputParam("test", std::make_tuple(info, testData));
   // Pass Labels.
   SetInputParam("labels", std::move(labels));
 
@@ -305,11 +305,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingModelCategoricalReuseTest)
   size_t testSize = testData.n_cols;
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   mlpackMain();
 
@@ -372,11 +372,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingMinSamplesTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("min_samples", 10);
   SetInputParam("confidence", 0.25);
@@ -404,11 +404,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingMinSamplesTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("min_samples", 2000);
   SetInputParam("confidence", 0.25);
@@ -441,11 +441,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingMaxSamplesTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("max_samples", 50000);
   SetInputParam("confidence", 0.95);
@@ -473,11 +473,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingMaxSamplesTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("max_samples", 5);
   SetInputParam("confidence", 0.95);
@@ -509,11 +509,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingConfidenceTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("confidence", 0.95);
 
@@ -540,11 +540,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingConfidenceTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   // Model with low confidence.
   SetInputParam("confidence", 0.25);
@@ -575,11 +575,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingPassesTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("passes", 1);
 
@@ -606,11 +606,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingPassesTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   // Model with larger number of passes.
   SetInputParam("passes", 100);
@@ -641,11 +641,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingBinarySplittingStrategyTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("numeric_split_strategy", (string) "binary");
   SetInputParam("max_samples", 50);
@@ -679,11 +679,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingDomingosSplittingStrategyTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("numeric_split_strategy", (string) "domingos");
   SetInputParam("max_samples", 50);
@@ -714,11 +714,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingDomingosSplittingStrategyTest)
     BOOST_FAIL("Cannot load test dataset vc2.csv!");
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, inputData)));
+  SetInputParam("training", std::make_tuple(info, inputData));
   SetInputParam("labels", std::move(labels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("numeric_split_strategy", (string) "domingos");
   SetInputParam("max_samples", 50);
@@ -756,11 +756,11 @@ BOOST_AUTO_TEST_CASE(HoeffdingBinningTest)
   modLabels = labels.cols(0, 49);
 
   // Input training data.
-  SetInputParam("training", std::move(std::make_tuple(info, modData)));
+  SetInputParam("training", std::make_tuple(info, modData));
   SetInputParam("labels", std::move(modLabels));
 
   // Input test data.
-  SetInputParam("test", std::move(std::make_tuple(info, testData)));
+  SetInputParam("test", std::make_tuple(info, testData));
 
   SetInputParam("numeric_split_strategy", (string) "domingos");
   SetInputParam("min_samples", 10);

src/mlpack/tests/main_tests/random_forest_test.cpp

@@ -186,6 +186,26 @@ BOOST_AUTO_TEST_CASE(RandomForestMinimumLeafSizeTest)
   Log::Fatal.ignoreInput = false;
 }
 
+/**
+ * Make sure maximum depth specified is always a positive number.
+ */
+BOOST_AUTO_TEST_CASE(RandomForestMaximumDepthTest)
+{
+  arma::mat inputData;
+  if (!data::Load("vc2.csv", inputData))
+    BOOST_FAIL("Cannot load train dataset vc2.csv!");
+
+  arma::Row<size_t> labels;
+  if (!data::Load("vc2_labels.txt", labels))
+    BOOST_FAIL("Cannot load labels for vc2_labels.txt");
+
+  SetInputParam("maximum_depth", (int) -1); // Invalid.
+
+  Log::Fatal.ignoreInput = true;
+  BOOST_REQUIRE_THROW(mlpackMain(), std::runtime_error);
+  Log::Fatal.ignoreInput = false;
+}
+
 /**
  * Make sure only one of training data or pre-trained model is passed.
  */
@@ -229,7 +249,7 @@ BOOST_AUTO_TEST_CASE(RandomForestDiffMinLeafSizeTest)
     BOOST_FAIL("Cannot load labels for vc2_labels.txt");
 
   bool success = false;
-  for (size_t trial = 0; trial < 3; ++trial)
+  for (size_t trial = 0; trial < 5; ++trial)
   {
     // Input training data.
     SetInputParam("training", inputData);
@@ -354,8 +374,8 @@ BOOST_AUTO_TEST_CASE(RandomForestDiffNumTreeTest)
     // Train for num_trees 10.
 
     // Input training data.
-    SetInputParam("training", std::move(inputData));
-    SetInputParam("labels", std::move(labels));
+    SetInputParam("training", inputData);
+    SetInputParam("labels", labels);
     SetInputParam("num_trees", (int) 10);
     SetInputParam("minimum_leaf_size", (int) 1);
 

src/mlpack/tests/q_learning_test.cpp

@@ -102,6 +102,74 @@ BOOST_AUTO_TEST_CASE(CartPoleWithDQN)
   BOOST_REQUIRE(converged);
 }
 
+//! 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));
+  model.Add<Linear<>>(4, 128);
+  model.Add<ReLULayer<>>();
+  model.Add<Linear<>>(128, 128);
+  model.Add<ReLULayer<>>();
+  model.Add<Linear<>>(128, 2);
+
+  // Set up the policy and replay method.
+  GreedyPolicy<CartPole> policy(1.0, 1000, 0.1);
+  PrioritizedReplay<CartPole> replayMethod(10, 10000, 0.6);
+
+  TrainingConfig config;
+  config.StepSize() = 0.01;
+  config.Discount() = 0.9;
+  config.TargetNetworkSyncInterval() = 100;
+  config.ExplorationSteps() = 100;
+  config.DoubleQLearning() = false;
+  config.StepLimit() = 200;
+
+  // Set up DQN agent.
+  QLearning<CartPole, decltype(model), AdamUpdate, decltype(policy),
+      decltype(replayMethod)>
+      agent(std::move(config), std::move(model), std::move(policy),
+          std::move(replayMethod));
+
+  arma::running_stat<double> averageReturn;
+  size_t episodes = 0;
+  bool converged = true;
+  while (true)
+  {
+    double episodeReturn = agent.Episode();
+    averageReturn(episodeReturn);
+    episodes += 1;
+
+    if (episodes > 1000)
+    {
+      Log::Debug << "Cart Pole with DQN failed." << std::endl;
+      converged = false;
+      break;
+    }
+
+    /**
+     * Reaching running average return 35 is enough to show it works.
+     * For the speed of the test case, I didn't set high criterion.
+     */
+    Log::Debug << "Average return: " << averageReturn.mean()
+        << " Episode return: " << episodeReturn << std::endl;
+    if (averageReturn.mean() > 35)
+    {
+      agent.Deterministic() = true;
+      arma::running_stat<double> testReturn;
+      for (size_t i = 0; i < 10; ++i)
+        testReturn(agent.Episode());
+
+      Log::Debug << "Average return in deterministic test: "
+          << testReturn.mean() << std::endl;
+      break;
+    }
+  }
+
+  BOOST_REQUIRE(converged);
+}
+
 //! Test Double DQN in Cart Pole task.
 BOOST_AUTO_TEST_CASE(CartPoleWithDoubleDQN)
 {

src/mlpack/tests/random_forest_test.cpp

@@ -223,7 +223,7 @@ BOOST_AUTO_TEST_CASE(UnweightedCategoricalLearningTest)
 
   // Train a random forest and a decision tree.
   RandomForest<> rf(trainingData, di, trainingLabels, 5, 25 /* 25 trees */, 1,
-      1e-7, MultipleRandomDimensionSelect(4));
+      1e-7, 0, MultipleRandomDimensionSelect(4));
   DecisionTree<> dt(trainingData, di, trainingLabels, 5, 5);
 
   // Get performance statistics on test data.
@@ -281,7 +281,7 @@ BOOST_AUTO_TEST_CASE(WeightedCategoricalLearningTest)
 
   // Build a random forest and a decision tree.
   RandomForest<> rf(fullData, di, fullLabels, 5, weights, 25 /* 25 trees */, 1,
-      1e-7, MultipleRandomDimensionSelect(4));
+      1e-7, 0, MultipleRandomDimensionSelect(4));
   DecisionTree<> dt(fullData, di, fullLabels, 5, weights, 5);
 
   // Get performance statistics on test data.
@@ -445,14 +445,14 @@ BOOST_AUTO_TEST_CASE(RandomForestCategoricalTrainReturnEntropy)
   // Test random forest on unweighted categorical dataset.
   RandomForest<> rf;
   double entropy = rf.Train(fullData, di, fullLabels, 5, 15 /* 15 trees */, 1,
-      1e-7, MultipleRandomDimensionSelect(3));
+      1e-7, 0, MultipleRandomDimensionSelect(3));
 
   BOOST_REQUIRE_EQUAL(std::isfinite(entropy), true);
 
   // Test random forest on weighted categorical dataset.
   RandomForest<> wrf;
   entropy = wrf.Train(fullData, di, fullLabels, 5, weights, 15 /* 15 trees */,
-      1, 1e-7, MultipleRandomDimensionSelect(3));
+      1, 1e-7, 0, MultipleRandomDimensionSelect(3));
 
   BOOST_REQUIRE_EQUAL(std::isfinite(entropy), true);
 }

src/mlpack/tests/reward_clipping_test.cpp

@@ -47,7 +47,7 @@ BOOST_AUTO_TEST_CASE(ClippedRewardTest)
 
   RewardClipping<Pendulum>::State state = rewardClipping.InitialSample();
   RewardClipping<Pendulum>::Action action;
-  action.action[0] = mlpack::math::Random(-1.0, 1.0);
+  action.action = mlpack::math::Random(-1.0, 1.0);
   double reward = rewardClipping.Sample(state, action);
 
   BOOST_REQUIRE(reward <= 2.0);

src/mlpack/tests/rl_components_test.cpp

@@ -16,6 +16,8 @@
 #include <mlpack/methods/reinforcement_learning/environment/mountain_car.hpp>
 #include <mlpack/methods/reinforcement_learning/environment/continuous_mountain_car.hpp>
 #include <mlpack/methods/reinforcement_learning/environment/cart_pole.hpp>
+#include <mlpack/methods/reinforcement_learning/environment/multiple_pole_cart.hpp>
+#include <mlpack/methods/reinforcement_learning/environment/continuous_multiple_pole_cart.hpp>
 #include <mlpack/methods/reinforcement_learning/environment/acrobot.hpp>
 #include <mlpack/methods/reinforcement_learning/environment/pendulum.hpp>
 #include <mlpack/methods/reinforcement_learning/replay/random_replay.hpp>
@@ -35,16 +37,25 @@ BOOST_AUTO_TEST_SUITE(RLComponentsTest)
  */
 BOOST_AUTO_TEST_CASE(SimplePendulumTest)
 {
-  const Pendulum task = Pendulum();
+  Pendulum task = Pendulum();
+  task.MaxSteps() = 5;
 
   Pendulum::State state = task.InitialSample();
   Pendulum::Action action;
-  action.action[0] = math::Random(-2.0, 2.0);
+  action.action = math::Random(-2.0, 2.0);
   double reward = task.Sample(state, action);
 
   // The reward is always negative. Check if not lower than lowest possible.
   BOOST_REQUIRE(reward >= -(pow(M_PI, 2) + 6.404));
 
+  BOOST_REQUIRE(!task.IsTerminal(state));
+
+  while (!task.IsTerminal(state))
+    task.Sample(state, action, state);
+
+  // Check if the number of steps performed is the same as the maximum allowed.
+  BOOST_REQUIRE_EQUAL(task.StepsPerformed(), 5);
+
   // The action is simply the torque. Check if dimension is 1.
   BOOST_REQUIRE_EQUAL(1, action.size);
 }
@@ -55,15 +66,24 @@ BOOST_AUTO_TEST_CASE(SimplePendulumTest)
  */
 BOOST_AUTO_TEST_CASE(SimpleContinuousMountainCarTest)
 {
-  const ContinuousMountainCar task = ContinuousMountainCar();
+  ContinuousMountainCar task = ContinuousMountainCar();
+  task.MaxSteps() = 5;
 
   ContinuousMountainCar::State state = task.InitialSample();
   ContinuousMountainCar::Action action;
-  action.action[0] = math::Random(-1.0, 1.0);
+  action.action = math::Random(-1.0, 1.0);
   double reward = task.Sample(state, action);
   // Maximum reward possible is 100.
   BOOST_REQUIRE(reward <= 100.0);
   BOOST_REQUIRE(!task.IsTerminal(state));
+
+  while (!task.IsTerminal(state))
+    task.Sample(state, action, state);
+
+  // Check if the number of steps performed is the same as the maximum allowed.
+  BOOST_REQUIRE_EQUAL(task.StepsPerformed(), 5);
+
+  // Check if the size of the action space is 1.
   BOOST_REQUIRE_EQUAL(1, action.size);
 }
 
@@ -73,7 +93,8 @@ BOOST_AUTO_TEST_CASE(SimpleContinuousMountainCarTest)
  */
 BOOST_AUTO_TEST_CASE(SimpleAcrobotTest)
 {
-  const Acrobot task = Acrobot();
+  Acrobot task = Acrobot();
+  task.MaxSteps() = 5;
 
   Acrobot::State state = task.InitialSample();
   Acrobot::Action action = Acrobot::Action::negativeTorque;
@@ -81,6 +102,14 @@ BOOST_AUTO_TEST_CASE(SimpleAcrobotTest)
 
   BOOST_REQUIRE_EQUAL(reward, -1.0);
   BOOST_REQUIRE(!task.IsTerminal(state));
+
+  while (!task.IsTerminal(state))
+    task.Sample(state, action, state);
+
+  // Check if the number of steps performed is the same as the maximum allowed.
+  BOOST_REQUIRE_EQUAL(task.StepsPerformed(), 5);
+
+  // Check if the size of the action space is 3.
   BOOST_REQUIRE_EQUAL(3, Acrobot::Action::size);
 }
 
@@ -90,7 +119,8 @@ BOOST_AUTO_TEST_CASE(SimpleAcrobotTest)
  */
 BOOST_AUTO_TEST_CASE(SimpleMountainCarTest)
 {
-  const MountainCar task = MountainCar();
+  MountainCar task = MountainCar();
+  task.MaxSteps() = 5;
 
   MountainCar::State state = task.InitialSample();
   MountainCar::Action action = MountainCar::Action::backward;
@@ -98,16 +128,25 @@ BOOST_AUTO_TEST_CASE(SimpleMountainCarTest)
 
   BOOST_REQUIRE_EQUAL(reward, -1.0);
   BOOST_REQUIRE(!task.IsTerminal(state));
+
+  while (!task.IsTerminal(state))
+    task.Sample(state, action, state);
+
+  // Check if the number of steps performed is the same as the maximum allowed.
+  BOOST_REQUIRE_EQUAL(task.StepsPerformed(), 5);
+
+  // Check if the size of the action space is 3.
   BOOST_REQUIRE_EQUAL(3, MountainCar::Action::size);
 }
 
 /**
- * Constructs a CartPole instance and check if the main rountine works as
+ * Constructs a CartPole instance and check if the main routine works as
  * it should be.
  */
 BOOST_AUTO_TEST_CASE(SimpleCartPoleTest)
 {
-  const CartPole task = CartPole();
+  CartPole task = CartPole();
+  task.MaxSteps() = 5;
 
   CartPole::State state = task.InitialSample();
   CartPole::Action action = CartPole::Action::backward;
@@ -115,9 +154,72 @@ BOOST_AUTO_TEST_CASE(SimpleCartPoleTest)
 
   BOOST_REQUIRE_EQUAL(reward, 1.0);
   BOOST_REQUIRE(!task.IsTerminal(state));
+
+  while (!task.IsTerminal(state))
+    task.Sample(state, action, state);
+
+  // Check if the number of steps performed is the same as the maximum allowed.
+  BOOST_REQUIRE_EQUAL(task.StepsPerformed(), 5);
+
   BOOST_REQUIRE_EQUAL(2, CartPole::Action::size);
 }
 
+/**
+ * Constructs a MultiplePoleCart instance and check if the main routine works as
+ * it should be.
+ */
+BOOST_AUTO_TEST_CASE(MultiplePoleCartTest)
+{
+  arma::vec poleLengths = {1, 0.5};
+  arma::vec poleMasses = {1, 1};
+  MultiplePoleCart task = MultiplePoleCart(2, poleLengths, poleMasses);
+  task.MaxSteps() = 5;
+
+  MultiplePoleCart::State state = task.InitialSample();
+  MultiplePoleCart::Action action = MultiplePoleCart::Action::backward;
+  double reward = task.Sample(state, action);
+
+  BOOST_REQUIRE_EQUAL(reward, 1.0);
+  BOOST_REQUIRE(!task.IsTerminal(state));
+
+  while (!task.IsTerminal(state))
+    task.Sample(state, action, state);
+
+  // Check if the number of steps performed is the same as the maximum allowed.
+  BOOST_REQUIRE_EQUAL(task.StepsPerformed(), 5);
+  BOOST_REQUIRE_EQUAL(2, MultiplePoleCart::Action::size);
+}
+
+/**
+ * Constructs a ContinuousMultiplePoleCart instance and check if the main 
+ * routine works as it should be.
+ */
+BOOST_AUTO_TEST_CASE(ContinuousMultiplePoleCartTest)
+{
+  arma::arma_rng::set_seed_random();
+
+  arma::vec poleLengths = {1, 0.5};
+  arma::vec poleMasses = {1, 1};
+  ContinuousMultiplePoleCart task = ContinuousMultiplePoleCart(2, poleLengths,
+      poleMasses);
+  task.MaxSteps() = 5;
+
+  ContinuousMultiplePoleCart::State state = task.InitialSample();
+  ContinuousMultiplePoleCart::Action action;
+  action.action = math::Random(-1.0, 1.0);
+  double reward = task.Sample(state, action);
+
+  BOOST_REQUIRE_EQUAL(reward, 1.0);
+  BOOST_REQUIRE(!task.IsTerminal(state));
+
+  while (!task.IsTerminal(state))
+    task.Sample(state, action, state);
+
+  // Check if the number of steps performed is the same as the maximum allowed.
+  BOOST_REQUIRE_EQUAL(task.StepsPerformed(), 5);
+  BOOST_REQUIRE_EQUAL(1, action.size);
+}
+
 /**
  * Construct a random replay instance and check if it works as
  * it should be.

src/mlpack/tests/sumtree_test.cpp

@@ -0,0 +1,94 @@
+/**
+ * @file sumtree_test.hpp
+ * @author Xiaohong
+ *
+ * Test for Sumtree implementation.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#include <mlpack/core.hpp>
+
+#include <mlpack/methods/reinforcement_learning/replay/sumtree.hpp>
+
+#include <boost/test/unit_test.hpp>
+#include "test_tools.hpp"
+
+using namespace mlpack;
+using namespace mlpack::rl;
+
+BOOST_AUTO_TEST_SUITE(SumTreeTest);
+
+/**
+ * Test that we set the element.
+ */
+BOOST_AUTO_TEST_CASE(SetElement)
+{
+  SumTree<double> sumtree(4);
+  sumtree.Set(0, 1.0);
+  sumtree.Set(1, 0.8);
+  sumtree.Set(2, 0.6);
+  sumtree.Set(3, 0.4);
+
+  BOOST_CHECK_CLOSE(sumtree.Sum(), 2.8, 1e-8);
+  BOOST_CHECK_CLOSE(sumtree.Sum(0, 1), 1.0, 1e-8);
+  BOOST_CHECK_CLOSE(sumtree.Sum(0, 3), 2.4, 1e-8);
+  BOOST_CHECK_CLOSE(sumtree.Sum(1, 4), 1.8, 1e-8);
+}
+
+/**
+ * Test that we get the element.
+ */
+BOOST_AUTO_TEST_CASE(GetElement)
+{
+  SumTree<double> sumtree(4);
+  sumtree.Set(0, 1.0);
+  sumtree.Set(1, 0.8);
+  sumtree.Set(2, 0.6);
+  sumtree.Set(3, 0.4);
+
+  BOOST_CHECK_CLOSE(sumtree.Get(0), 1.0, 1e-8);
+  BOOST_CHECK_CLOSE(sumtree.Get(1), 0.8, 1e-8);
+  BOOST_CHECK_CLOSE(sumtree.Get(2), 0.6, 1e-8);
+  BOOST_CHECK_CLOSE(sumtree.Get(3), 0.4, 1e-8);
+}
+
+/**
+ * Test that we find the highest index in the array such that
+ * Sum(arr[0] + arr[1] + arr[2] ... + arr[i]) <= mass.
+ */
+BOOST_AUTO_TEST_CASE(FindPrefixSum)
+{
+  SumTree<double> sumtree(4);
+  sumtree.Set(0, 1.0);
+  sumtree.Set(1, 0.8);
+  sumtree.Set(2, 0.6);
+  sumtree.Set(3, 0.4);
+
+  BOOST_CHECK_EQUAL(sumtree.FindPrefixSum(0), 0);
+  BOOST_CHECK_EQUAL(sumtree.FindPrefixSum(1), 1);
+  BOOST_CHECK_EQUAL(sumtree.FindPrefixSum(2.8), 3);
+  BOOST_CHECK_EQUAL(sumtree.FindPrefixSum(3.0), 3);
+}
+
+/**
+ * Test that we find the highest index in the array such that
+ * sum(arr[0] + arr[1] + arr[2] ... + arr[i]) <= mass.
+ */
+BOOST_AUTO_TEST_CASE(BatchUpdate)
+{
+  SumTree<double> sumtree(4);
+  arma::ucolvec indices = {0, 1, 2, 3};
+  arma::colvec data = {1.0, 0.8, 0.6, 0.4};
+
+  sumtree.BatchUpdate(indices, data);
+
+  BOOST_CHECK_EQUAL(sumtree.FindPrefixSum(0), 0);
+  BOOST_CHECK_EQUAL(sumtree.FindPrefixSum(1), 1);
+  BOOST_CHECK_EQUAL(sumtree.FindPrefixSum(2.8), 3);
+  BOOST_CHECK_EQUAL(sumtree.FindPrefixSum(3.0), 3);
+}
+
+BOOST_AUTO_TEST_SUITE_END();