FirstStateNextStateGenerator.h

Go to the documentation of this file.

#pragma once
 
#include <set>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
 
#include "storm-pomdp/beliefs/abstraction/NoAbstraction.h"
 
#include "storm-pomdp/beliefs/storage/BeliefBuilder.h"
#include "storm-pomdp/beliefs/utility/types.h"
 
#include "storm/utility/constants.h"
 
namespace storm::pomdp::beliefs {
 
// Forward declare handle
namespace detail {
template<typename PomdpType, typename BeliefType, typename PreAbstractionType, typename PostAbstractionType, typename DiscoverCallback>
struct NextStateGeneratorHandle;
}

template<typename PomdpType, typename BeliefType>
class FirstStateNextStateGenerator {
   public:
    template<typename PreAbstractionType, typename PostAbstractionType, typename DiscoverCallbackType>
    using Handle = detail::NextStateGeneratorHandle<PomdpType, BeliefType, PreAbstractionType, PostAbstractionType, DiscoverCallbackType>;

    FirstStateNextStateGenerator(PomdpType const& pomdp);

    void setRewardModel(std::string const& rewardModelName = "");
    bool hasRewardModel() const;
    void unsetRewardModel();

    BeliefType computeInitialBelief() const;

    uint64_t getBeliefNumberOfActions(BeliefType const& belief) const;

    typename PomdpType::ValueType getBeliefActionReward(BeliefType const& belief, uint64_t const& localActionIndex) const;

    std::set<std::string> getBeliefActionChoiceLabels(BeliefType const& belief, uint64_t const& localActionIndex) const;

    template<typename DiscoverCallbackType>
    auto getHandle(DiscoverCallbackType& discoverCallback) {
        return Handle<NoAbstractionType const, NoAbstractionType const, DiscoverCallbackType>{this->pomdp, NoAbstraction, NoAbstraction, discoverCallback};
    }

    template<typename PreAbstractionType, typename DiscoverCallbackType>
    auto getPreAbstractionHandle(PreAbstractionType& preAbstraction, DiscoverCallbackType& discoverCallback) {
        return Handle<PreAbstractionType, NoAbstractionType const, DiscoverCallbackType>{this->pomdp, preAbstraction, NoAbstraction, discoverCallback};
    }

    template<typename PostAbstractionType, typename DiscoverCallbackType>
    auto getPostAbstractionHandle(PostAbstractionType& postAbstraction, DiscoverCallbackType& discoverCallback) {
        return Handle<NoAbstractionType const, PostAbstractionType, DiscoverCallbackType>{this->pomdp, NoAbstraction, postAbstraction, discoverCallback};
    }

    template<typename PreAbstractionType, typename PostAbstractionType, typename DiscoverCallbackType>
    auto getPrePostAbstractionHandle(PreAbstractionType& preAbstraction, PostAbstractionType& postAbstraction, DiscoverCallbackType& discoverCallback) {
        return Handle<PreAbstractionType, PostAbstractionType, DiscoverCallbackType>{this->pomdp, preAbstraction, postAbstraction, discoverCallback};
    }
 
   private:
    PomdpType const& pomdp;
    std::vector<typename PomdpType::ValueType> actionRewards;
};
 
namespace detail {

template<typename PomdpType, typename BeliefType, typename PreAbstractionType, typename PostAbstractionType, typename DiscoverCallback>
struct NextStateGeneratorHandle {
    using BeliefValueType = typename BeliefType::ValueType;
 
    template<typename... CallBackArgs>
    void operator()(BeliefType const& belief, uint64_t localActionIndex, CallBackArgs const&... additionalCallbackArgs) {
        applyPreAbstraction(belief, localActionIndex, std::forward<CallBackArgs const>(additionalCallbackArgs)...);
    }
 
    PomdpType const& pomdp;
    PreAbstractionType& preAbstraction;
    PostAbstractionType& postAbstraction;
    DiscoverCallback& discoverCallback;
 
   private:
    template<typename... CallBackArgs>
    void applyPreAbstraction(BeliefType const& belief, uint64_t localActionIndex, CallBackArgs const&... additionalCallbackArgs) {
        if constexpr (isNoAbstraction<PreAbstractionType>) {
            computeSuccessorBeliefs(belief, localActionIndex, storm::utility::one<BeliefValueType>(), DefaultActionObservation, additionalCallbackArgs...);
        } else {
            preAbstraction.abstract(
                belief, localActionIndex,
                [this, &localActionIndex, &additionalCallbackArgs...](
                    BeliefType&& preBel, BeliefValueType&& preVal, BeliefActionObservationType actionObservation, auto const&... additionalPreAbstractionArgs) {
                    this->computeSuccessorBeliefs(preBel, localActionIndex, std::move(preVal), actionObservation,
                                                  std::forward<CallBackArgs const>(additionalCallbackArgs)...,
                                                  std::forward<decltype(additionalPreAbstractionArgs)>(additionalPreAbstractionArgs)...);
                });
        }
    }

    std::unordered_map<BeliefObservationType, BeliefValueType> computeSuccessorObservations(BeliefType const& belief, uint64_t localActionIndex,
                                                                                            BeliefActionObservationType actionObservation) {
        std::unordered_map<BeliefObservationType, BeliefValueType> successorObservations;
        belief.forEach([&localActionIndex, &successorObservations, &actionObservation, this](BeliefStateType const& state, BeliefValueType const& beliefValue) {
            for (auto const& pomdpTransition : pomdp.getTransitionMatrix().getRow(state, localActionIndex)) {
                if (!storm::utility::isZero(pomdpTransition.getValue())) {
                    auto const obs = (pomdp.getNrObservations() * actionObservation) + pomdp.getObservation(pomdpTransition.getColumn());
                    STORM_LOG_ASSERT(obs < InvalidObservation, "Encoded belief observation " << obs << " exceeds the valid observation range.");
                    BeliefValueType const val = beliefValue * storm::utility::convertNumber<BeliefValueType>(pomdpTransition.getValue());
                    if (auto [insertionIt, inserted] = successorObservations.emplace(static_cast<BeliefObservationType>(obs), val); !inserted) {
                        insertionIt->second += val;
                    }
                }
            }
        });
 
        // Adjust the distribution to diminish numerical inaccuracies a bit
        if constexpr (!storm::NumberTraits<BeliefValueType>::IsExact || !storm::NumberTraits<typename PomdpType::ValueType>::IsExact) {
            if (successorObservations.size() == 1) {
                successorObservations.begin()->second = storm::utility::one<BeliefValueType>();
            }
        }
        return successorObservations;
    }
 
    template<typename... CallBackArgs>
    void computeSuccessorBeliefs(BeliefType const& belief, uint64_t localActionIndex, BeliefValueType const& transitionProbability,
                                 BeliefActionObservationType actionObservation, CallBackArgs const&... additionalCallbackArgs) {
        // For each successor observation we build the successor belief
        auto const successorObservations = computeSuccessorObservations(belief, localActionIndex, actionObservation);
        for (auto const& successorObsValue : successorObservations) {
            BeliefBuilder<BeliefType> builder;
            builder.setObservation(successorObsValue.first);
            belief.forEach([&builder, &localActionIndex, &successorObsValue, this](BeliefStateType const& state, BeliefValueType const& beliefValue) {
                for (auto const& pomdpTransition : pomdp.getTransitionMatrix().getRow(state, localActionIndex)) {
                    if (!storm::utility::isZero(pomdpTransition.getValue()) &&
                        pomdp.getObservation(pomdpTransition.getColumn()) == (successorObsValue.first % pomdp.getNrObservations())) {
                        BeliefValueType const prob =
                            beliefValue * storm::utility::convertNumber<BeliefValueType>(pomdpTransition.getValue()) / successorObsValue.second;
                        STORM_LOG_ASSERT(prob > storm::utility::zero<BeliefValueType>(), "Invalid belief probability " << prob << ".");
                        builder.addValue(pomdpTransition.getColumn(), prob);
                    }
                }
            });
            if constexpr (!storm::NumberTraits<BeliefValueType>::IsExact || !storm::NumberTraits<typename PomdpType::ValueType>::IsExact) {
                if (successorObservations.size() == 1) {
                    // The observation probability was forced to one, so normalize the successor belief separately.
                    builder.normalize();
                }
            }
            applyPostAbstraction(builder.build(), static_cast<BeliefValueType>(successorObsValue.second * transitionProbability),
                                 std::forward<CallBackArgs const>(additionalCallbackArgs)...);
        }
    }
 
    template<typename... CallBackArgs>
    void applyPostAbstraction(BeliefType&& belief, BeliefValueType&& transitionProbability, CallBackArgs const&... additionalCallbackArgs) {
        if constexpr (isNoAbstraction<PostAbstractionType>) {
            discoverCallback(std::move(belief), std::move(transitionProbability), additionalCallbackArgs...);
        } else {
            postAbstraction.abstract(
                std::move(belief), std::move(transitionProbability),
                [this, &additionalCallbackArgs...](BeliefType&& postBel, BeliefValueType&& postVal, auto&&... additionalPostAbstractionArgs) {
                    discoverCallback(std::move(postBel), std::move(postVal), std::forward<CallBackArgs const>(additionalCallbackArgs)...,
                                     std::forward<decltype(additionalPostAbstractionArgs)>(additionalPostAbstractionArgs)...);
                });
        }
    }
};
 
}  // namespace detail
}  // namespace storm::pomdp::beliefs