one_step_sarsa_worker.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_RL_WORKER_ONE_STEP_SARSA_WORKER_HPP
#define MLPACK_METHODS_RL_WORKER_ONE_STEP_SARSA_WORKER_HPP

#include <mlpack/methods/reinforcement_learning/training_config.hpp>

namespace mlpack {
namespace rl {

template <
  typename EnvironmentType,
  typename NetworkType,
  typename UpdaterType,
  typename PolicyType
>
class OneStepSarsaWorker
{
 public:
  using StateType = typename EnvironmentType::State;
  using ActionType = typename EnvironmentType::Action;
  using TransitionType = std::tuple<StateType, ActionType, double, StateType,
      ActionType>;

  OneStepSarsaWorker(
      const UpdaterType& updater,
      const EnvironmentType& environment,
      const TrainingConfig& config,
      bool deterministic):
      updater(updater),
      environment(environment),
      config(config),
      deterministic(deterministic),
      pending(config.UpdateInterval())
  { Reset(); }

  void Initialize(NetworkType& learningNetwork)
  {
    updater.Initialize(learningNetwork.Parameters().n_rows,
        learningNetwork.Parameters().n_cols);
    // Build local network.
    network = learningNetwork;
  }

  bool Step(NetworkType& learningNetwork,
            NetworkType& targetNetwork,
            size_t& totalSteps,
            PolicyType& policy,
            double& totalReward)
  {
    // Interact with the environment.
    if (action == ActionType::size)
    {
      // Invalid action means we are at the beginning of an episode.
      arma::colvec actionValue;
      network.Predict(state.Encode(), actionValue);
      action = policy.Sample(actionValue, deterministic);
    }
    StateType nextState;
    double reward = environment.Sample(state, action, nextState);
    bool terminal = environment.IsTerminal(nextState);
    arma::colvec actionValue;
    network.Predict(nextState.Encode(), actionValue);
    ActionType nextAction = policy.Sample(actionValue, deterministic);

    episodeReturn += reward;
    steps++;

    terminal = terminal || steps >= config.StepLimit();
    if (deterministic)
    {
      if (terminal)
      {
        totalReward = episodeReturn;
        Reset();
        // Sync with latest learning network.
        network = learningNetwork;
        return true;
      }
      state = nextState;
      action = nextAction;
      return false;
    }

    #pragma omp atomic
    totalSteps++;

    pending[pendingIndex++] =
        std::make_tuple(state, action, reward, nextState, nextAction);

    if (terminal || pendingIndex >= config.UpdateInterval())
    {
      // Initialize the gradient storage.
      arma::mat totalGradients(learningNetwork.Parameters().n_rows,
          learningNetwork.Parameters().n_cols, arma::fill::zeros);
      for (size_t i = 0; i < pending.size(); ++i)
      {
        TransitionType &transition = pending[i];

        // Compute the target state-action value.
        arma::colvec actionValue;
        #pragma omp critical
        {
          targetNetwork.Predict(
              std::get<3>(transition).Encode(), actionValue);
        };
        double targetActionValue = 0;
        if (!(terminal && i == pending.size() - 1))
          targetActionValue = actionValue[std::get<4>(transition)];
        targetActionValue = std::get<2>(transition) +
            config.Discount() * targetActionValue;

        // Compute the training target for current state.
        network.Forward(std::get<0>(transition).Encode(), actionValue);
        actionValue[std::get<1>(transition)] = targetActionValue;

        // Compute gradient.
        arma::mat gradients;
        network.Backward(actionValue, gradients);

        // Accumulate gradients.
        totalGradients += gradients;
      }

      // Clamp the accumulated gradients.
      totalGradients.transform(
          [&](double gradient)
          { return std::min(std::max(gradient, -config.GradientLimit()),
          config.GradientLimit()); });

      // Perform async update of the global network.
      updater.Update(learningNetwork.Parameters(),
          config.StepSize(), totalGradients);

      // Sync the local network with the global network.
      network = learningNetwork;

      pendingIndex = 0;
    }

    // Update global target network.
    if (totalSteps % config.TargetNetworkSyncInterval() == 0)
    {
      #pragma omp critical
      { targetNetwork = learningNetwork; }
    }

    policy.Anneal();

    if (terminal)
    {
      totalReward = episodeReturn;
      Reset();
      return true;
    }
    state = nextState;
    action = nextAction;
    return false;
  }

 private:
  void Reset()
  {
    steps = 0;
    episodeReturn = 0;
    pendingIndex = 0;
    state = environment.InitialSample();
    action = ActionType::size;
  }

  UpdaterType updater;

  EnvironmentType environment;

  TrainingConfig config;

  bool deterministic;

  size_t steps;

  double episodeReturn;

  std::vector<TransitionType> pending;

  size_t pendingIndex;

  NetworkType network;

  StateType state;

  ActionType action;
};

} // namespace rl
} // namespace mlpack

#endif