prioritized_replay.hpp

Go to the documentation of this file.

#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 {

template <typename EnvironmentType>
class PrioritizedReplay
{
 public:
  using ActionType = typename EnvironmentType::Action;

  using StateType = typename EnvironmentType::State;

  PrioritizedReplay()
  { /* Nothing to do here. */ }

  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);
  }

  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;
    }
  }

  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;
  }

  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();
  }

  void UpdatePriorities(arma::ucolvec& indices, arma::colvec& priorities)
  {
      arma::colvec alphaPri = alpha * priorities;
      maxPriority = std::max(maxPriority, arma::max(priorities));
      idxSum.BatchUpdate(indices, alphaPri);
  }

  const size_t& Size()
  {
    return full ? capacity : position;
  }

  void BetaAnneal()
  {
    beta = beta + (1 - initialBeta) * 1.0 / replayBetaIters;
  }

  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:
  size_t batchSize;

  size_t capacity;

  size_t position;

  arma::mat states;

  arma::icolvec actions;

  arma::colvec rewards;

  arma::mat nextStates;

  arma::icolvec isTerminal;

  bool full;

  double alpha;

  double maxPriority;

  double initialBeta;

  double beta;

  size_t replayBetaIters;

  SumTree<double> idxSum;

  arma::ucolvec sampledIndices;

  arma::rowvec weights;
};

} // namespace rl
} // namespace mlpack

#endif