diff --git a/HISTORY.md b/HISTORY.md index 67d39452d2..d427e2f4f1 100644 --- a/HISTORY.md +++ b/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). diff --git a/src/mlpack/methods/reinforcement_learning/environment/cart_pole.hpp b/src/mlpack/methods/reinforcement_learning/environment/cart_pole.hpp index 3a30a522fb..b3fe48d375 100644 --- a/src/mlpack/methods/reinforcement_learning/environment/cart_pole.hpp +++ b/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. diff --git a/src/mlpack/methods/reinforcement_learning/q_learning.hpp b/src/mlpack/methods/reinforcement_learning/q_learning.hpp index 4ec62dc0fb..afbdcff33e 100644 --- a/src/mlpack/methods/reinforcement_learning/q_learning.hpp +++ b/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. */ diff --git a/src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp b/src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp index 12c5d744fc..f6ef0078a1 100644 --- a/src/mlpack/methods/reinforcement_learning/q_learning_impl.hpp +++ b/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 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(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 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::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::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::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; } diff --git a/src/mlpack/methods/reinforcement_learning/q_networks/CMakeLists.txt b/src/mlpack/methods/reinforcement_learning/q_networks/CMakeLists.txt index 01a3dcebac..eb5d2014b8 100644 --- a/src/mlpack/methods/reinforcement_learning/q_networks/CMakeLists.txt +++ b/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. diff --git a/src/mlpack/methods/reinforcement_learning/q_networks/categorical_dqn.hpp b/src/mlpack/methods/reinforcement_learning/q_networks/categorical_dqn.hpp new file mode 100644 index 0000000000..b52110d744 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#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 +> +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(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*> + (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 noisyLayerIndex; + + //! Locally-stored softmax activation function. + Softmax<> softMax; + + //! Locally-stored activations from softMax. + arma::mat activations; +}; + +} // namespace rl +} // namespace mlpack + +#endif diff --git a/src/mlpack/methods/reinforcement_learning/training_config.hpp b/src/mlpack/methods/reinforcement_learning/training_config.hpp index b2e9fc7001..a9dbd8544d 100644 --- a/src/mlpack/methods/reinforcement_learning/training_config.hpp +++ b/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 diff --git a/src/mlpack/tests/q_learning_test.cpp b/src/mlpack/tests/q_learning_test.cpp index 0a19d65f04..308ed2b37f 100644 --- a/src/mlpack/tests/q_learning_test.cpp +++ b/src/mlpack/tests/q_learning_test.cpp @@ -17,9 +17,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -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 policy(1.0, 1000, 0.1, 0.99); + RandomReplay replayMethod(32, 4000); + + TrainingConfig config; + config.IsCategorical() = true; + config.ExplorationSteps() = 32; + + // Set up the module. Note that we use a custom network here. + FFN, GaussianInitialization> module( + EmptyLoss<>(), GaussianInitialization(0, 0.1)); + module.Add>(4, 128); + module.Add>(); + module.Add>(128, 2 * config.AtomSize()); + + // Adding the module to the CategoricalDQN network. + CategoricalDQN<> network(module, config); + + // Set up DQN agent. + QLearning + agent(config, network, policy, replayMethod); + + converged = testAgent(agent, 40, 1000, 20); + if (converged) + break; + } + BOOST_REQUIRE(converged); +} + BOOST_AUTO_TEST_SUITE_END();