Merge branch 'master' of https://github.com/mlpack/mlpack into FixLinearLayer

HISTORY.md

@@ -1,5 +1,7 @@
 ### mlpack ?.?.?
 ###### ????-??-??
+  * Added Categorical DQN to q_networks (#2454).
+
   * Added N-step DQN to q_networks (#2461).
 
   * Add Silhoutte Score metric and Pairwise Distances (#2406).

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

@@ -123,7 +123,7 @@ 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 = 1.0) :
       maxSteps(maxSteps),
       gravity(gravity),
       massCart(massCart),
@@ -177,8 +177,6 @@ class CartPole
     // 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.

src/mlpack/methods/reinforcement_learning/q_learning.hpp

@@ -94,6 +94,11 @@ class QLearning
    */
   void TrainAgent();
 
+  /**
+   * Trains the DQN agent of categorical type.
+   */
+  void TrainCategoricalAgent();
+
   /**
    * Select an action, given an agent.
    */

src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp

@@ -48,9 +48,12 @@ QLearning<
     totalSteps(0),
     deterministic(false)
 {
+  // To copy over the network structure.
+  targetNetwork = learningNetwork;
+
   // Set up q-learning network.
-  if (learningNetwork.Parameters().is_empty())
-    learningNetwork.ResetParameters();
+  learningNetwork.ResetParameters();
+  targetNetwork.ResetParameters();
 
   #if ENS_VERSION_MAJOR == 1
   this->updater.Initialize(learningNetwork.Parameters().n_rows,
@@ -62,7 +65,8 @@ QLearning<
                                    learningNetwork.Parameters().n_cols);
   #endif
 
-  targetNetwork = learningNetwork;
+  // Initialize the target network with the parameters of learning network.
+  targetNetwork.Parameters() = learningNetwork.Parameters();
 }
 
 template <
@@ -192,7 +196,118 @@ void QLearning<
   }
   // Update target network.
   if (totalSteps % config.TargetNetworkSyncInterval() == 0)
-    targetNetwork = learningNetwork;
+    targetNetwork.Parameters() = learningNetwork.Parameters();
+
+  if (totalSteps > config.ExplorationSteps())
+    policy.Anneal();
+}
+
+template <
+  typename EnvironmentType,
+  typename NetworkType,
+  typename UpdaterType,
+  typename BehaviorPolicyType,
+  typename ReplayType
+>
+void QLearning<
+  EnvironmentType,
+  NetworkType,
+  UpdaterType,
+  BehaviorPolicyType,
+  ReplayType
+>::TrainCategoricalAgent()
+{
+  // Start experience replay.
+
+  // Sample from previous experience.
+  arma::mat sampledStates;
+  std::vector<ActionType> sampledActions;
+  arma::colvec sampledRewards;
+  arma::mat sampledNextStates;
+  arma::icolvec isTerminal;
+
+  replayMethod.Sample(sampledStates, sampledActions, sampledRewards,
+      sampledNextStates, isTerminal);
+
+  size_t atomSize = config.AtomSize();
+  arma::rowvec support = arma::linspace<arma::rowvec>(config.VMin(),
+      config.VMax(), atomSize);
+
+  size_t batchSize = sampledNextStates.n_cols;
+
+  // Compute action value for next state with target network.
+  arma::mat nextActionValues;
+  targetNetwork.Predict(sampledNextStates, nextActionValues);
+
+  arma::Col<size_t> nextAction;
+  if (config.DoubleQLearning())
+  {
+    // If use double Q-Learning, use learning network to select the best action.
+    arma::mat nextActionValues;
+    learningNetwork.Predict(sampledNextStates, nextActionValues);
+    nextAction = BestAction(nextActionValues);
+  }
+  else
+  {
+    nextAction = BestAction(nextActionValues);
+  }
+
+  arma::mat nextDists, nextDist(atomSize, batchSize);
+  targetNetwork.Forward(sampledNextStates, nextDists);
+  for (size_t i = 0; i < batchSize; ++i)
+  {
+    nextDist.col(i) = nextDists(nextAction(i) * atomSize, i,
+        arma::size(atomSize, 1));
+  }
+
+  arma::mat tZ = (arma::conv_to<arma::mat>::from(config.Discount() *
+      ((1 - isTerminal) * support)).each_col() + sampledRewards).t();
+  tZ = arma::clamp(tZ, config.VMin(), config.VMax());
+  arma::mat b = (tZ - config.VMin()) / (config.VMax() - config.VMin()) *
+      (atomSize - 1);
+  arma::mat l = arma::floor(b);
+  arma::mat u = arma::ceil(b);
+
+  arma::mat projDistUpper = nextDist % (u - b);
+  arma::mat projDistLower = nextDist % (b - l);
+
+  arma::mat projDist = arma::zeros<arma::mat>(arma::size(nextDist));
+  for (size_t batchNo = 0; batchNo < batchSize; batchNo++)
+  {
+    for (size_t j = 0; j < atomSize; j++)
+    {
+      projDist(l(j, batchNo), batchNo) += projDistUpper(j, batchNo);
+      projDist(u(j, batchNo), batchNo) += projDistLower(j, batchNo);
+    }
+  }
+  arma::mat dists;
+  learningNetwork.Forward(sampledStates, dists);
+  arma::mat lossGradients = arma::zeros<arma::mat>(arma::size(dists));
+  for (size_t i = 0; i < batchSize; ++i)
+  {
+    lossGradients(sampledActions[i].action * atomSize, i,
+        arma::size(atomSize, 1)) = -(projDist.col(i) / (1e-10 + dists(
+        sampledActions[i].action * atomSize, i, arma::size(atomSize, 1))));
+  }
+  // Learn from experience.
+  arma::mat gradients;
+  learningNetwork.Backward(sampledStates, lossGradients, gradients);
+
+  #if ENS_VERSION_MAJOR == 1
+  updater.Update(learningNetwork.Parameters(), config.StepSize(), gradients);
+  #else
+  updatePolicy->Update(learningNetwork.Parameters(), config.StepSize(),
+      gradients);
+  #endif
+
+  if (config.NoisyQLearning() == true)
+  {
+    learningNetwork.ResetNoise();
+    targetNetwork.ResetNoise();
+  }
+  // Update target network.
+  if (totalSteps % config.TargetNetworkSyncInterval() == 0)
+    targetNetwork.Parameters() = learningNetwork.Parameters();
 
   if (totalSteps > config.ExplorationSteps())
     policy.Anneal();
@@ -262,7 +377,10 @@ double QLearning<
 
     if (deterministic || totalSteps < config.ExplorationSteps())
       continue;
-    TrainAgent();
+    if (config.IsCategorical())
+      TrainCategoricalAgent();
+    else
+      TrainAgent();
   }
   return totalReturn;
 }

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

@@ -3,6 +3,7 @@
 set(SOURCES
   simple_dqn.hpp
   dueling_dqn.hpp
+  categorical_dqn.hpp
 )
 
 # Add directory name to sources.

src/mlpack/methods/reinforcement_learning/q_networks/categorical_dqn.hpp

@@ -0,0 +1,246 @@
+/**
+ * @file methods/reinforcement_learning/q_networks/categorical_dqn.hpp
+ * @author Nishant Kumar
+ *
+ * This file contains the implementation of the categorical deep q network.
+ *
+ * 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_CATEGORICAL_DQN_HPP
+#define MLPACK_METHODS_RL_CATEGORICAL_DQN_HPP
+
+#include <mlpack/prereqs.hpp>
+#include <mlpack/methods/ann/ffn.hpp>
+#include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
+#include <mlpack/methods/ann/layer/layer.hpp>
+#include <mlpack/methods/ann/loss_functions/mean_squared_error.hpp>
+#include <mlpack/methods/ann/loss_functions/empty_loss.hpp>
+#include "../training_config.hpp"
+
+namespace mlpack {
+namespace rl {
+
+using namespace mlpack::ann;
+
+/**
+ * Implementation of the Categorical Deep Q-Learning network.
+ * For more information, see the following.
+ *
+ * @code
+ * @misc{bellemare2017distributional,
+ *   author  = {Marc G. Bellemare, Will Dabney, Rémi Munos},
+ *   title   = {A Distributional Perspective on Reinforcement Learning},
+ *   year    = {2017},
+ *   url     = {http://arxiv.org/abs/1707.06887}
+ * }
+ * @endcode
+ * 
+ * @tparam OutputLayerType The output layer type of the network.
+ * @tparam InitType The initialization type used for the network.
+ * @tparam NetworkType The type of network used for simple dqn.
+ */
+template<
+  typename OutputLayerType = EmptyLoss<>,
+  typename InitType = GaussianInitialization,
+  typename NetworkType = FFN<OutputLayerType, InitType>
+>
+class CategoricalDQN
+{
+ public:
+  /**
+   * Default constructor.
+   */
+  CategoricalDQN() : network(), isNoisy(false)
+  { /* Nothing to do here. */ }
+
+  /**
+   * Construct an instance of CategoricalDQN class.
+   *
+   * @param inputDim Number of inputs.
+   * @param h1 Number of neurons in hiddenlayer-1.
+   * @param h2 Number of neurons in hiddenlayer-2.
+   * @param outputDim Number of neurons in output layer.
+   * @param config Hyper-parameters for categorical dqn.
+   * @param isNoisy Specifies whether the network needs to be of type noisy.
+   * @param init Specifies the initialization rule for the network.
+   * @param outputLayer Specifies the output layer type for network.
+   */
+  CategoricalDQN(const int inputDim,
+                 const int h1,
+                 const int h2,
+                 const int outputDim,
+                 TrainingConfig config,
+                 const bool isNoisy = false,
+                 InitType init = InitType(),
+                 OutputLayerType outputLayer = OutputLayerType()):
+      network(outputLayer, init),
+      atomSize(config.AtomSize()),
+      vMin(config.VMin()),
+      vMax(config.VMax()),
+      isNoisy(isNoisy)
+  {
+    network.Add(new Linear<>(inputDim, h1));
+    network.Add(new ReLULayer<>());
+    if (isNoisy)
+    {
+      noisyLayerIndex.push_back(network.Model().size());
+      network.Add(new NoisyLinear<>(h1, h2));
+      network.Add(new ReLULayer<>());
+      noisyLayerIndex.push_back(network.Model().size());
+      network.Add(new NoisyLinear<>(h2, outputDim * atomSize));
+    }
+    else
+    {
+      network.Add(new Linear<>(h1, h2));
+      network.Add(new ReLULayer<>());
+      network.Add(new Linear<>(h2, outputDim * atomSize));
+    }
+  }
+
+  /**
+   * Construct an instance of CategoricalDQN class from a pre-constructed network.
+   *
+   * @param network The network to be used by CategoricalDQN class.
+   * @param config Hyper-parameters for categorical dqn.
+   * @param isNoisy Specifies whether the network needs to be of type noisy.
+   */
+  CategoricalDQN(NetworkType& network,
+                 TrainingConfig config,
+                 const bool isNoisy = false):
+      network(std::move(network)),
+      atomSize(config.AtomSize()),
+      vMin(config.VMin()),
+      vMax(config.VMax()),
+      isNoisy(isNoisy)
+  { /* Nothing to do here. */ }
+
+  /**
+   * Predict the responses to a given set of predictors. The responses will
+   * reflect the output of the given output layer as returned by the
+   * output layer function.
+   *
+   * If you want to pass in a parameter and discard the original parameter
+   * object, be sure to use std::move to avoid unnecessary copy.
+   *
+   * @param state Input state.
+   * @param actionValue Matrix to put output action values of states input.
+   */
+  void Predict(const arma::mat state, arma::mat& actionValue)
+  {
+    arma::mat q_atoms;
+    network.Predict(state, q_atoms);
+    activations.copy_size(q_atoms);
+    actionValue.set_size(q_atoms.n_rows / atomSize, q_atoms.n_cols);
+    arma::rowvec support = arma::linspace<arma::rowvec>(vMin, vMax, atomSize);
+    for (size_t i = 0; i < q_atoms.n_rows; i += atomSize)
+    {
+      arma::mat activation = activations.rows(i, i + atomSize - 1);
+      arma::mat input = q_atoms.rows(i, i + atomSize - 1);
+      softMax.Forward(input, activation);
+      activations.rows(i, i + atomSize - 1) = activation;
+      actionValue.row(i/atomSize) = support * activation;
+    }
+  }
+
+  /**
+   * Perform the forward pass of the states in real batch mode.
+   *
+   * @param state The input state.
+   * @param dist The predicted distributions.
+   */
+  void Forward(const arma::mat state, arma::mat& dist)
+  {
+    arma::mat q_atoms;
+    network.Forward(state, q_atoms);
+    activations.copy_size(q_atoms);
+    for (size_t i = 0; i < q_atoms.n_rows; i += atomSize)
+    {
+      arma::mat activation = activations.rows(i, i + atomSize - 1);
+      arma::mat input = q_atoms.rows(i, i + atomSize - 1);
+      softMax.Forward(input, activation);
+      activations.rows(i, i + atomSize - 1) = activation;
+    }
+    dist = activations;
+  }
+
+  /**
+   * Resets the parameters of the network.
+   */
+  void ResetParameters()
+  {
+    network.ResetParameters();
+  }
+
+  /**
+   * Resets noise of the network, if the network is of type noisy.
+   */
+  void ResetNoise()
+  {
+    for (size_t i = 0; i < noisyLayerIndex.size(); i++)
+    {
+      boost::get<NoisyLinear<>*>
+          (network.Model()[noisyLayerIndex[i]])->ResetNoise();
+    }
+  }
+
+  //! Return the Parameters.
+  const arma::mat& Parameters() const { return network.Parameters(); }
+  //! Modify the Parameters.
+  arma::mat& Parameters() { return network.Parameters(); }
+
+  /**
+   * Perform the backward pass of the state in real batch mode.
+   *
+   * @param state The input state.
+   * @param lossGradients The loss gradients.
+   * @param gradient The gradient.
+   */
+  void Backward(const arma::mat state,
+                arma::mat& lossGradients,
+                arma::mat& gradient)
+  {
+    arma::mat activationGradients(arma::size(activations));
+    for (size_t i = 0; i < activations.n_rows; i += atomSize)
+    {
+      arma::mat activationGrad;
+      arma::mat lossGrad = lossGradients.rows(i, i + atomSize - 1);
+      arma::mat activation = activations.rows(i, i + atomSize - 1);
+      softMax.Backward(activation, lossGrad, activationGrad);
+      activationGradients.rows(i, i + atomSize - 1) = activationGrad;
+    }
+    network.Backward(state, activationGradients, gradient);
+  }
+
+ private:
+  //! Locally-stored network.
+  NetworkType network;
+
+  //! Locally-stored number of atoms.
+  size_t atomSize;
+
+  //! Locally-stored minimum value of support.
+  double vMin;
+
+  //! Locally-stored maximum value of support.
+  double vMax;
+
+  //! Locally-stored check for noisy network.
+  bool isNoisy;
+
+  //! Locally-stored indexes of noisy layers in the network.
+  std::vector<size_t> noisyLayerIndex;
+
+  //! Locally-stored softmax activation function.
+  Softmax<> softMax;
+
+  //! Locally-stored activations from softMax.
+  arma::mat activations;
+};
+
+} // namespace rl
+} // namespace mlpack
+
+#endif

src/mlpack/methods/reinforcement_learning/training_config.hpp

@@ -29,7 +29,11 @@ class TrainingConfig
       discount(0.99),
       gradientLimit(40),
       doubleQLearning(false),
-      noisyQLearning(false)
+      noisyQLearning(false),
+      isCategorical(false),
+      atomSize(51),
+      vMin(0),
+      vMax(200)
   { /* Nothing to do here. */ }
 
   TrainingConfig(
@@ -42,7 +46,11 @@ class TrainingConfig
       double discount,
       double gradientLimit,
       bool doubleQLearning,
-      bool noisyQLearning) :
+      bool noisyQLearning,
+      bool isCategorical,
+      size_t atomSize,
+      double vMin,
+      double vMax) :
       numWorkers(numWorkers),
       updateInterval(updateInterval),
       targetNetworkSyncInterval(targetNetworkSyncInterval),
@@ -52,7 +60,11 @@ class TrainingConfig
       discount(discount),
       gradientLimit(gradientLimit),
       doubleQLearning(doubleQLearning),
-      noisyQLearning(noisyQLearning)
+      noisyQLearning(noisyQLearning),
+      isCategorical(isCategorical),
+      atomSize(atomSize),
+      vMin(vMin),
+      vMax(vMax)
   { /* Nothing to do here. */ }
 
   //! Get the amount of workers.
@@ -109,6 +121,26 @@ class TrainingConfig
   //! Modify the indicator of double q-learning.
   bool& NoisyQLearning() { return noisyQLearning; }
 
+  //! Get the indicator of categorical q-learning.
+  bool IsCategorical() const { return isCategorical; }
+  //! Modify the indicator of categorical q-learning.
+  bool& IsCategorical() { return isCategorical; }
+
+  //! Get the number of atoms.
+  size_t AtomSize() const { return atomSize; }
+  //! Modify the number of atoms.
+  size_t& AtomSize() { return atomSize; }
+
+  //! Get the minimum value for support.
+  double VMin() const { return vMin; }
+  //! Modify the minimum value for support.
+  double& VMin() { return vMin; }
+
+  //! Get the maximum value for support.
+  double VMax() const { return vMax; }
+  //! Modify the maximum value for support.
+  double& VMax() { return vMax; }
+
  private:
   /**
    * Locally-stored number of workers.
@@ -172,6 +204,30 @@ class TrainingConfig
    * This is valid only for q-learning agent.
    */
   bool noisyQLearning;
+
+  /**
+   * Locally-stored indicator for categorical q-learning.
+   * This is valid only for q-learning agent.
+   */
+  bool isCategorical;
+
+  /**
+   * Locally-stored number of atoms to be used.
+   * This is valid only for categorical q-network.
+   */
+  size_t atomSize;
+
+  /**
+   * Locally-stored minimum value of support.
+   * This is valid only for categorical q-network.
+   */
+  double vMin;
+
+  /**
+   * Locally-stored maximum value of support.
+   * This is valid only for categorical q-network.
+   */
+  double vMax;
 };
 
 } // namespace rl

src/mlpack/tests/q_learning_test.cpp

@@ -17,9 +17,11 @@
 #include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
 #include <mlpack/methods/ann/layer/layer.hpp>
 #include <mlpack/methods/ann/loss_functions/mean_squared_error.hpp>
+#include <mlpack/methods/ann/loss_functions/empty_loss.hpp>
 #include <mlpack/methods/reinforcement_learning/q_learning.hpp>
 #include <mlpack/methods/reinforcement_learning/q_networks/simple_dqn.hpp>
 #include <mlpack/methods/reinforcement_learning/q_networks/dueling_dqn.hpp>
+#include <mlpack/methods/reinforcement_learning/q_networks/categorical_dqn.hpp>
 #include <mlpack/methods/reinforcement_learning/environment/mountain_car.hpp>
 #include <mlpack/methods/reinforcement_learning/environment/acrobot.hpp>
 #include <mlpack/methods/reinforcement_learning/environment/cart_pole.hpp>
@@ -320,7 +322,6 @@ BOOST_AUTO_TEST_CASE(CartPoleWithDuelingDQN)
   BOOST_REQUIRE(converged);
 }
 
-
 //! Test Dueling DQN in Cart Pole task with Prioritized Replay.
 BOOST_AUTO_TEST_CASE(CartPoleWithDuelingDQNPrioritizedReplay)
 {
@@ -376,7 +377,6 @@ BOOST_AUTO_TEST_CASE(CartPoleWithNoisyDQN)
   BOOST_REQUIRE(converged);
 }
 
-
 //! Test Dueling-Double-Noisy DQN in Cart Pole task.
 BOOST_AUTO_TEST_CASE(CartPoleWithDuelingDoubleNoisyDQN)
 {
@@ -465,4 +465,43 @@ BOOST_AUTO_TEST_CASE(CartPoleWithNStepPrioritizedDQN)
   BOOST_REQUIRE(converged);
 }
 
+//! Test Categorical DQN in Cart Pole task.
+BOOST_AUTO_TEST_CASE(CartPoleWithCategoricalDQN)
+{
+  // It isn't guaranteed that the network will converge in the specified number
+  // of iterations.
+  bool converged = false;
+  for (size_t trial = 0; trial < 3; ++trial)
+  {
+    Log::Debug << "Trial number: " << trial << std::endl;
+
+    // Set up the policy and replay method.
+    GreedyPolicy<CartPole> policy(1.0, 1000, 0.1, 0.99);
+    RandomReplay<CartPole> replayMethod(32, 4000);
+
+    TrainingConfig config;
+    config.IsCategorical() = true;
+    config.ExplorationSteps() = 32;
+
+    // Set up the module. Note that we use a custom network here.
+    FFN<EmptyLoss<>, GaussianInitialization> module(
+        EmptyLoss<>(), GaussianInitialization(0, 0.1));
+    module.Add<Linear<>>(4, 128);
+    module.Add<ReLULayer<>>();
+    module.Add<Linear<>>(128, 2 * config.AtomSize());
+
+    // Adding the module to the CategoricalDQN network.
+    CategoricalDQN<> network(module, config);
+
+    // Set up DQN agent.
+    QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy)>
+        agent(config, network, policy, replayMethod);
+
+    converged = testAgent<decltype(agent)>(agent, 40, 1000, 20);
+    if (converged)
+      break;
+  }
+  BOOST_REQUIRE(converged);
+}
+
 BOOST_AUTO_TEST_SUITE_END();