simple_dqn.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_RL_SIMPLE_DQN_HPP
#define MLPACK_METHODS_RL_SIMPLE_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>

namespace mlpack {
namespace rl {

template<
  typename OutputLayerType = ann::MeanSquaredError,
  typename InitType = ann::GaussianInitialization,
  typename NetworkType = ann::FFN<OutputLayerType, InitType>
>
class SimpleDQN
{
 public:
  SimpleDQN() : network(), isNoisy(false)
  { /* Nothing to do here. */ }

  SimpleDQN(const int h1,
            const int h2,
            const int outputDim,
            const bool isNoisy = false,
            InitType init = InitType(),
            OutputLayerType outputLayer = OutputLayerType()):
      network(outputLayer, init),
      isNoisy(isNoisy)
  {
    network.Add(new ann::Linear(h1));
    network.Add(new ann::ReLU());
    if (isNoisy)
    {
      noisyLayerIndex.push_back(network.Network().size());
      network.Add(new ann::NoisyLinear(h2));
      network.Add(new ann::ReLU());
      noisyLayerIndex.push_back(network.Network().size());
      network.Add(new ann::NoisyLinear(outputDim));
    }
    else
    {
      network.Add(new ann::Linear(h2));
      network.Add(new ann::ReLU());
      network.Add(new ann::Linear(outputDim));
    }
  }

  SimpleDQN(NetworkType& network, const bool isNoisy = false):
      network(network),
      isNoisy(isNoisy)
  { /* Nothing to do here. */ }

  void Predict(const arma::mat state, arma::mat& actionValue)
  {
    network.Predict(state, actionValue);
  }

  void Forward(const arma::mat state, arma::mat& target)
  {
    network.Forward(state, target);
  }

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

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

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

  void Backward(const arma::mat state, arma::mat& target, arma::mat& gradient)
  {
    network.Backward(state, target, gradient);
  }

 private:
  NetworkType network;

  bool isNoisy;

  std::vector<size_t> noisyLayerIndex;
};

} // namespace rl
} // namespace mlpack

#endif