POMCP.hpp

Go to the documentation of this file.

#ifndef AI_TOOLBOX_POMDP_POMCP_HEADER_FILE
#define AI_TOOLBOX_POMDP_POMCP_HEADER_FILE
 
#include <unordered_map>
 
#include <AIToolbox/Logging.hpp>
#include <AIToolbox/Seeder.hpp>
#include <AIToolbox/Utils/Probability.hpp>
#include <AIToolbox/POMDP/Types.hpp>
#include <AIToolbox/POMDP/TypeTraits.hpp>
#include <AIToolbox/MDP/Algorithms/Utils/Rollout.hpp>
 
namespace AIToolbox::POMDP {
    template <IsGenerativeModel M>
    class POMCP {
        public:
            using SampleBelief = std::vector<size_t>;
 
            struct BeliefNode;
            using BeliefNodes = std::unordered_map<size_t, BeliefNode>;
 
            struct ActionNode {
                BeliefNodes children;
                double V = 0.0;
                unsigned N = 0;
            };
            using ActionNodes = std::vector<ActionNode>;
 
            struct BeliefNode {
                BeliefNode() : N(0) {}
                BeliefNode(size_t s) : belief(1, s), N(0) {}
                ActionNodes children;
                SampleBelief belief;
                unsigned N;
            };
 
            POMCP(const M& m, size_t beliefSize, unsigned iterations, double exp);
 
            size_t sampleAction(const Belief& b, unsigned horizon);
 
            size_t sampleAction(size_t a, size_t o, unsigned horizon);
 
            void setBeliefSize(size_t beliefSize);
 
            void setIterations(unsigned iter);
 
            void setExploration(double exp);
 
            const M& getModel() const;
 
            const BeliefNode& getGraph() const;
 
            size_t getBeliefSize() const;
 
            unsigned getIterations() const;
 
            double getExploration() const;
 
        private:
            const M& model_;
            size_t S, A, beliefSize_;
            unsigned iterations_, maxDepth_;
            double exploration_;
 
            SampleBelief sampleBelief_;
            BeliefNode graph_;
 
            mutable RandomEngine rand_;
 
            size_t runSimulation(unsigned horizon);
 
            double simulate(BeliefNode & b, size_t s, unsigned horizon);
 
            template <typename Iterator>
            Iterator findBestA(Iterator begin, Iterator end);
 
            template <typename Iterator>
            Iterator findBestBonusA(Iterator begin, Iterator end, unsigned count);
 
            SampleBelief makeSampledBelief(const Belief & b);
    };
 
    template <IsGenerativeModel M>
    POMCP<M>::POMCP(const M& m, const size_t beliefSize, const unsigned iter, const double exp) :
            model_(m), S(model_.getS()), A(model_.getA()), beliefSize_(beliefSize),
            iterations_(iter), exploration_(exp), graph_(), rand_(Seeder::getSeed()) {}
 
    template <IsGenerativeModel M>
    size_t POMCP<M>::sampleAction(const Belief& b, const unsigned horizon) {
        // Reset graph
        graph_ = BeliefNode(A);
        graph_.children.resize(A);
        graph_.belief = makeSampledBelief(b);
 
        return runSimulation(horizon);
    }
 
    template <IsGenerativeModel M>
    size_t POMCP<M>::sampleAction(const size_t a, const size_t o, const unsigned horizon) {
        const auto & obs = graph_.children[a].children;
 
        auto it = obs.find(o);
        if ( it == obs.end() ) {
            AI_LOGGER(AI_SEVERITY_WARNING, "Observation " << o << " never experienced in simulation, restarting with uniform belief..");
            auto b = Belief(S); b.fill(1.0/S);
            return sampleAction(b, horizon);
        }
 
        // Here we need an additional step, because *it is contained by graph_.
        // If we just move assign, graph_ is first going to delete everything it
        // contains (included *it), and then we are going to move unallocated memory
        // into graph_! So we move *it outside of the graph_ hierarchy, so that
        // we can then assign safely.
        { auto tmp = std::move(it->second); graph_ = std::move(tmp); }
 
        if ( ! graph_.belief.size() ) {
            AI_LOGGER(AI_SEVERITY_WARNING, "POMCP lost track of the belief, restarting with uniform..");
            auto b = Belief(S); b.fill(1.0/S);
            return sampleAction(b, horizon);
        }
 
        // We resize here in case we didn't have time to sample the new
        // head node. In this case, the new head may not have children.
        // This would break the UCT call.
        graph_.children.resize(A);
 
        return runSimulation(horizon);
    }
 
    template <IsGenerativeModel M>
    size_t POMCP<M>::runSimulation(const unsigned horizon) {
        if ( !horizon ) return 0;
 
        maxDepth_ = horizon;
        std::uniform_int_distribution<size_t> generator(0, graph_.belief.size()-1);
 
        for (unsigned i = 0; i < iterations_; ++i )
            simulate(graph_, graph_.belief.at(generator(rand_)), 0);
 
        auto begin = std::begin(graph_.children);
        return std::distance(begin, findBestA(begin, std::end(graph_.children)));
    }
 
    template <IsGenerativeModel M>
    double POMCP<M>::simulate(BeliefNode & b, const size_t s, const unsigned depth) {
        b.N++;
 
        auto begin = std::begin(b.children);
        const size_t a = std::distance(begin, findBestBonusA(begin, std::end(b.children), b.N));
 
        auto [s1, o, rew] = model_.sampleSOR(s, a);
 
        auto & aNode = b.children[a];
 
        {
            double futureRew = 0.0;
            // We need to append the node anyway to perform the belief
            // update for the next timestep.
            auto ot = aNode.children.find(o);
            if ( ot == std::end(aNode.children) ) {
                aNode.children.emplace(std::piecewise_construct,
                                       std::forward_as_tuple(o),
                                       std::forward_as_tuple(s1));
                // This stops automatically if we go out of depth
                futureRew = rollout(model_, s1, maxDepth_ - depth + 1, rand_);
            }
            else {
                ot->second.belief.push_back(s1);
                // We only go deeper if needed (maxDepth_ is always at least 1).
                if ( depth + 1 < maxDepth_ && !model_.isTerminal(s1) ) {
                    // Since most memory is allocated on the leaves,
                    // we do not allocate on node creation but only when
                    // we are actually descending into a node. If the node
                    // already has memory this should not do anything in
                    // any case.
                    ot->second.children.resize(A);
                    futureRew = simulate( ot->second, s1, depth + 1 );
                }
            }
 
            rew += model_.getDiscount() * futureRew;
        }
 
        // Action update
        aNode.N++;
        aNode.V += ( rew - aNode.V ) / static_cast<double>(aNode.N);
 
        return rew;
    }
 
    template <IsGenerativeModel M>
    template <typename Iterator>
    Iterator POMCP<M>::findBestA(Iterator begin, Iterator end) {
        return std::max_element(begin, end, [](const ActionNode & lhs, const ActionNode & rhs){ return lhs.V < rhs.V; });
    }
 
    template <IsGenerativeModel M>
    template <typename Iterator>
    Iterator POMCP<M>::findBestBonusA(Iterator begin, Iterator end, const unsigned count) {
        // Count here can be as low as 1.
        // Since log(1) = 0, and 0/0 = error, we add 1.0.
        const double logCount = std::log(count + 1.0);
        // We use this function to produce a score for each action. This can be easily
        // substituted with something else to produce different POMCP variants.
        auto evaluationFunction = [this, logCount](const ActionNode & an){
                return an.V + exploration_ * std::sqrt( logCount / an.N );
        };
 
        auto bestIterator = begin++;
        double bestValue = evaluationFunction(*bestIterator);
 
        for ( ; begin < end; ++begin ) {
            const double actionValue = evaluationFunction(*begin);
            if ( actionValue > bestValue ) {
                bestValue = actionValue;
                bestIterator = begin;
            }
        }
 
        return bestIterator;
    }
 
    template <IsGenerativeModel M>
    typename POMCP<M>::SampleBelief POMCP<M>::makeSampledBelief(const Belief & b) {
        SampleBelief belief;
        belief.reserve(beliefSize_);
 
        for ( size_t i = 0; i < beliefSize_; ++i )
            belief.push_back(sampleProbability(S, b, rand_));
 
        return belief;
    }
 
    template <IsGenerativeModel M>
    void POMCP<M>::setBeliefSize(const size_t beliefSize) {
        beliefSize_ = beliefSize;
    }
 
    template <IsGenerativeModel M>
    void POMCP<M>::setIterations(const unsigned iter) {
        iterations_ = iter;
    }
 
    template <IsGenerativeModel M>
    void POMCP<M>::setExploration(const double exp) {
        exploration_ = exp;
    }
 
    template <IsGenerativeModel M>
    const M& POMCP<M>::getModel() const {
        return model_;
    }
 
    template <IsGenerativeModel M>
    const typename POMCP<M>::BeliefNode& POMCP<M>::getGraph() const {
        return graph_;
    }
 
    template <IsGenerativeModel M>
    size_t POMCP<M>::getBeliefSize() const {
        return beliefSize_;
    }
 
    template <IsGenerativeModel M>
    unsigned POMCP<M>::getIterations() const {
        return iterations_;
    }
 
    template <IsGenerativeModel M>
    double POMCP<M>::getExploration() const {
        return exploration_;
    }
}
 
#endif