dueling_dqn.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_RL_DUELING_DQN_HPP
#define MLPACK_METHODS_RL_DUELING_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_types.hpp>
#include <mlpack/methods/ann/loss_functions/mean_squared_error.hpp>
#include <mlpack/methods/ann/loss_functions/empty_loss.hpp>

namespace mlpack {
namespace rl {

template <
  typename OutputLayerType = ann::EmptyLoss,
  typename InitType = ann::GaussianInitialization,
  typename CompleteNetworkType = ann::FFN<OutputLayerType, InitType>,
  typename FeatureNetworkType = ann::MultiLayer<arma::mat>,
  typename AdvantageNetworkType = ann::MultiLayer<arma::mat>,
  typename ValueNetworkType = ann::MultiLayer<arma::mat>
>
class DuelingDQN
{
 public:
  DuelingDQN() : isNoisy(false)
  {
    // TODO: this really ought to use a DAG network, but that's not implemented
    // yet.
    featureNetwork = new ann::MultiLayer<arma::mat>();
    valueNetwork = new ann::MultiLayer<arma::mat>();
    advantageNetwork = new ann::MultiLayer<arma::mat>();
    concat = new ann::Concat();

    concat->Add(valueNetwork);
    concat->Add(advantageNetwork);
    completeNetwork.Add(featureNetwork);
    completeNetwork.Add(concat);
  }

  DuelingDQN(const int h1,
             const int h2,
             const int outputDim,
             const bool isNoisy = false,
             InitType init = InitType(),
             OutputLayerType outputLayer = OutputLayerType()):
      completeNetwork(outputLayer, init),
      isNoisy(isNoisy)
  {
    featureNetwork = new ann::MultiLayer<arma::mat>();
    featureNetwork->Add(new ann::Linear(h1));
    featureNetwork->Add(new ann::ReLU());

    valueNetwork = new ann::MultiLayer<arma::mat>();
    advantageNetwork = new ann::MultiLayer<arma::mat>();

    if (isNoisy)
    {
      noisyLayerIndex.push_back(valueNetwork->Network().size());
      valueNetwork->Add(new ann::NoisyLinear(h2));
      advantageNetwork->Add(new ann::NoisyLinear(h2));

      valueNetwork->Add(new ann::ReLU());
      advantageNetwork->Add(new ann::ReLU());

      noisyLayerIndex.push_back(valueNetwork->Network().size());
      valueNetwork->Add(new ann::NoisyLinear(1));
      advantageNetwork->Add(new ann::NoisyLinear(outputDim));
    }
    else
    {
      valueNetwork->Add(new ann::Linear(h2));
      valueNetwork->Add(new ann::ReLU());
      valueNetwork->Add(new ann::Linear(1));

      advantageNetwork->Add(new ann::Linear(h2));
      advantageNetwork->Add(new ann::ReLU());
      advantageNetwork->Add(new ann::Linear(outputDim));
    }

    concat = new ann::Concat();
    concat->Add(valueNetwork);
    concat->Add(advantageNetwork);

    completeNetwork.Add(featureNetwork);
    completeNetwork.Add(concat);
  }

  DuelingDQN(FeatureNetworkType& featureNetwork,
             AdvantageNetworkType& advantageNetwork,
             ValueNetworkType& valueNetwork,
             const bool isNoisy = false):
      featureNetwork(featureNetwork),
      advantageNetwork(advantageNetwork),
      valueNetwork(valueNetwork),
      isNoisy(isNoisy)
  {
    concat = new ann::Concat();
    concat->Add(valueNetwork);
    concat->Add(advantageNetwork);
    completeNetwork.Add(featureNetwork);
    completeNetwork.Add(concat);
  }

  DuelingDQN(const DuelingDQN& /* model */) : isNoisy(false)
  { /* Nothing to do here. */ }

  void operator = (const DuelingDQN& model)
  {
    *valueNetwork = *model.valueNetwork;
    *advantageNetwork = *model.advantageNetwork;
    *featureNetwork = *model.featureNetwork;
    isNoisy = model.isNoisy;
    noisyLayerIndex = model.noisyLayerIndex;
  }

  void Predict(const arma::mat state, arma::mat& actionValue)
  {
    arma::mat advantage, value, networkOutput;
    completeNetwork.Predict(state, networkOutput);
    value = networkOutput.row(0);
    advantage = networkOutput.rows(1, networkOutput.n_rows - 1);
    actionValue = advantage.each_row() + (value - arma::mean(advantage));
  }

  void Forward(const arma::mat state, arma::mat& actionValue)
  {
    arma::mat advantage, value, networkOutput;
    completeNetwork.Forward(state, networkOutput);
    value = networkOutput.row(0);
    advantage = networkOutput.rows(1, networkOutput.n_rows - 1);
    actionValue = advantage.each_row() +
        (value - arma::mean(advantage));
    this->actionValues = actionValue;
  }

  void Backward(const arma::mat state, arma::mat& target, arma::mat& gradient)
  {
    arma::mat gradLoss;
    lossFunction.Backward(this->actionValues, target, gradLoss);

    arma::mat gradValue = arma::sum(gradLoss);
    arma::mat gradAdvantage = gradLoss.each_row() - arma::mean(gradLoss);

    arma::mat grad = arma::join_cols(gradValue, gradAdvantage);
    completeNetwork.Backward(state, grad, gradient);
  }

  void Reset(const size_t inputDimensionality = 0)
  {
    completeNetwork.Reset(inputDimensionality);
  }

  void ResetNoise()
  {
    for (size_t i = 0; i < noisyLayerIndex.size(); i++)
    {
      dynamic_cast<ann::NoisyLinear*>(
          (valueNetwork->Network()[noisyLayerIndex[i]]))->ResetNoise();
      dynamic_cast<ann::NoisyLinear*>(
          (advantageNetwork->Network()[noisyLayerIndex[i]]))->ResetNoise();
    }
  }

  const arma::mat& Parameters() const { return completeNetwork.Parameters(); }
  arma::mat& Parameters() { return completeNetwork.Parameters(); }

 private:
  CompleteNetworkType completeNetwork;

  ann::Concat* concat;

  FeatureNetworkType* featureNetwork;

  AdvantageNetworkType* advantageNetwork;

  ValueNetworkType* valueNetwork;

  bool isNoisy;

  std::vector<size_t> noisyLayerIndex;

  arma::mat actionValues;

  ann::MeanSquaredError lossFunction;
};

} // namespace rl
} // namespace mlpack

#endif