robust_values.hpp

#pragma once

#include "../RMDP.hpp"
#include "values.hpp"
#include <functional>
#include <type_traits>
#include "../cpp11-range-master/range.hpp"

namespace craam { namespace algorithms{

using namespace std;
using namespace util::lang;

// *******************************************************
// Nature definitions 
// *******************************************************

template<class T>
using NatureResponse = vec_scal_t (*)(numvec const& v, numvec const& p, T threshold);

template<class T>
using NatureInstance = pair<NatureResponse<T>, T>;


inline vec_scal_t robust_l1(const numvec& v, const numvec& p, prec_t threshold){
    assert(v.size() == p.size());
    return worstcase_l1(v,p,threshold);
}

inline vec_scal_t optimistic_l1(const numvec& v, const numvec& p, prec_t threshold){
    assert(v.size() == p.size());
    //TODO: this could be faster without copying the vector and just modifying the function
    numvec minusv(v.size());
    transform(begin(v), end(v), begin(minusv), negate<prec_t>());
    auto&& result = worstcase_l1(minusv,p,threshold);
    return make_pair(result.first, -result.second);
}

template<class T>
inline vec_scal_t robust_unbounded(const numvec& v, const numvec& p, T){
    assert(v.size() == p.size());
    numvec dist(v.size(),0.0);
    long index = min_element(begin(v), end(v)) - begin(v);
    dist[index] = 1;
    return make_pair(dist,v[index]);
}

template<class T>
inline vec_scal_t optimistic_unbounded(const numvec& v, const numvec& p, T){
    assert(v.size() == p.size());
    numvec dist(v.size(),0.0);
    long index = max_element(begin(v), end(v)) - begin(v);
    dist[index] = 1;
    return make_pair(dist,v[index]);
}

// *******************************************************
// RegularAction computation methods
// *******************************************************


template<class T>
inline vec_scal_t value_action(const RegularAction& action, const numvec& valuefunction,
                        prec_t discount, const NatureInstance<T>& nature){

    const numvec& rewards = action.get_outcome().get_rewards();
    const indvec& nonzero_indices = action.get_outcome().get_indices();

    numvec qvalues(rewards.size()); // values for individual states - used by nature.

    #pragma omp simd
    for(size_t i = 0; i < rewards.size(); i++){
        qvalues[i] = rewards[i] + discount * valuefunction[nonzero_indices[i]];
    }

    return nature.first(qvalues, action.get_outcome().get_probabilities(), nature.second);
}



// *******************************************************
// WeightedOutcomeAction computation methods
// *******************************************************

template<class T>
inline vec_scal_t value_action(const WeightedOutcomeAction& action, numvec const& valuefunction,
                                prec_t discount, const NatureInstance<T> nature) {

    assert(action.get_distribution().size() == action.get_outcomes().size());

    if(action.get_outcomes().empty())
        throw invalid_argument("Action with no action.get_outcomes().");

    numvec outcomevalues(action.size());
    for(size_t i = 0; i < action.size(); i++)
        outcomevalues[i] = action[i].value(valuefunction, discount);

    return nature.first(outcomevalues, action.get_distribution(), nature.second);
}


// *******************************************************
// State computation methods
// *******************************************************


template<class AType, class T>
inline vec_scal_t
value_fix_state(const SAState<AType>& state, numvec const& valuefunction, prec_t discount,
                              long actionid, const NatureInstance<T>& nature) {
   // this is the terminal state, return 0
    if(state.is_terminal()) return make_pair(numvec(0),0);

    assert(actionid >= 0 && actionid < long(state.size()));

    if(actionid < 0 || actionid >= (long) state.size()) throw range_error("invalid actionid: " + to_string(actionid) + " for action count: " + to_string(state.get_actions().size()) );

    const auto& action = state[actionid];
    // cannot assume that the action is valid
    if(!state.is_valid(actionid)) throw invalid_argument("Cannot take an invalid action");

    return value_action(action, valuefunction, discount, nature);
}



template<typename AType, typename T>
inline ind_vec_scal_t
value_max_state(const SAState<AType>& state, const numvec& valuefunction,
                prec_t discount, const NatureInstance<T>& nature) {

    if(state.is_terminal())
        return make_tuple(-1,numvec(),0);

    prec_t maxvalue = -numeric_limits<prec_t>::infinity();

    long result = -1;
    numvec result_outcome;

    for(size_t i = 0; i < state.get_actions().size(); i++){
        const auto& action = state[i];

        // skip invalid state.get_actions()
        if(!state.is_valid(i)) continue;

        auto value = value_action(action, valuefunction, discount, nature);
        if(value.second > maxvalue){
            maxvalue = value.second;
            result = i;
            result_outcome = move(value.first);
        }
    }

    // if the result has not been changed, that means that all actions are invalid
    if(result == -1)
        throw invalid_argument("all actions are invalid.");

    return make_tuple(result,result_outcome,maxvalue);
}


// **************************************************************************
// Helper classes to handle computing of the best response
// **************************************************************************

struct SolutionRobust : public Solution {
    vector<numvec> natpolicy;

    SolutionRobust() : Solution(), natpolicy(0) {};

    SolutionRobust(size_t statecount): Solution(statecount), natpolicy(statecount, numvec(0)) {};

    SolutionRobust(numvec valuefunction, indvec policy):
            Solution(move(valuefunction), move(policy)),
            natpolicy(this->valuefunction.size(), numvec(0)) {};

    SolutionRobust(numvec valuefunction, indvec policy,
             vector<numvec> natpolicy, prec_t residual = -1, long iterations = -1) :
        Solution(move(valuefunction), move(policy), residual, iterations),
        natpolicy(move(natpolicy)) {};
};
template<class T>
class PolicyNature : public PolicyDeterministic {
public:
    using solution_type = SolutionRobust;

    vector<NatureInstance<T>> natspec;

    PolicyNature(indvec policy, vector<NatureInstance<T>> natspec):
        PolicyDeterministic(move(policy)), natspec(move(natspec)) {};

    PolicyNature(vector<NatureInstance<T>> natspec):
        PolicyDeterministic(indvec(0)), natspec(move(natspec)) {};

    SolutionRobust new_solution(size_t statecount, numvec valuefunction) const {
        if(natspec.size() != statecount)
            throw invalid_argument("Size of nature specification does not match the number of states.");

        process_valuefunction(statecount, valuefunction);
        SolutionRobust solution =  SolutionRobust(move(valuefunction), process_policy(statecount));
        return solution;
    }

    template<class SType>
    prec_t update_solution(SolutionRobust& solution, const SType& state, long stateid,
                            const numvec& valuefunction, prec_t discount) const{

        prec_t newvalue = 0;
        // check whether this state should only be evaluated or also optimized
        if(policy.empty() || policy[stateid] < 0){    // optimizing
            tie(solution.policy[stateid], solution.natpolicy[stateid], newvalue) = value_max_state(state, valuefunction, discount, natspec[stateid]);
        }else{// fixed-action, do not copy
            prec_t newvalue;
            tie(solution.natpolicy[stateid], newvalue) = value_fix_state(state, valuefunction, discount, policy[stateid], natspec[stateid]);
        }
        return newvalue;
    }

    template<class SType>
    prec_t update_value(const SolutionRobust& solution, const SType& state, long stateid,
                            const numvec& valuefunction, prec_t discount) const{

        return value_fix_state(state, valuefunction, discount, solution.policy[stateid],
                solution.natpolicy[stateid]);
    }

};


template<class T>
PolicyNature<T> uniform_nature(size_t statecount, NatureResponse<T> nature,
                            T threshold){
    return PolicyNature<T>(vector<NatureInstance<T>>(statecount, make_pair(nature, threshold)));
}

template<class Model, class T>
PolicyNature<T> uniform_nature(const Model& m, NatureResponse<T> nature,
                            T threshold){
    return PolicyNature<T>(vector<NatureInstance<T>>(m.state_count(), make_pair(nature, threshold)));
}


// **************************************************************************
// Convenient interface methods
// **************************************************************************

namespace internal{

template <class T1, class T2>
vector<pair<T1,T2>> zip(const vector<T1>& v1, const vector<T2>& v2){
    assert(v1.size() == v2.size());
    vector<pair<T1,T2>> result(v1.size());
    for(size_t i=0; i< v1.size(); i++){
        result[i] = make_pair(v1[i], v2[i]);
    }
    return result;
}

template <class T1, class T2>
vector<pair<T1,T2>> zip(const T1& v1, const vector<T2>& v2){
    vector<pair<T1,T2>> result(v2.size());
    for(size_t i=0; i< v2.size(); i++){
        result[i] = make_pair(v1, v2[i]);
    }
    return result;
}
}

template<class SType, class T = prec_t >
inline auto rsolve_vi(const GRMDP<SType>& mdp, prec_t discount,
                        const vector<NatureResponse<T>>& nature, const vector<T>& thresholds,
                        numvec valuefunction=numvec(0), const indvec& policy = numvec(0),
                        unsigned long iterations=MAXITER, prec_t maxresidual=SOLPREC)
    {
    assert(nature.size() == thresholds.size());
    assert(nature.size() == mdp.state_count());

    return vi_gs<SType, PolicyNature<T>>(mdp, discount, move(valuefunction), 
            PolicyNature<T>(policy,internal::zip(nature,thresholds)), 
            iterations, maxresidual);
}

template<class SType, class T = prec_t >
inline auto rsolve_vi(const GRMDP<SType>& mdp, prec_t discount,
                        const NatureResponse<T>& nature, const vector<T>& thresholds,
                        numvec valuefunction=numvec(0), const indvec& policy = numvec(0),
                        unsigned long iterations=MAXITER, prec_t maxresidual=SOLPREC)
    {
    assert(nature.size() == thresholds.size());
    assert(nature.size() == mdp.state_count());

    return vi_gs<SType, PolicyNature<T>>(mdp, discount, move(valuefunction), 
            PolicyNature<T>(policy,internal::zip(nature,thresholds)), 
            iterations, maxresidual);
}



template<class SType, class T = prec_t>
inline auto rsolve_mpi(const GRMDP<SType>& mdp, prec_t discount,
                const vector<NatureResponse<T>>& nature, const vector<T>& thresholds,
                const numvec& valuefunction=numvec(0), const indvec& policy = indvec(0),
                unsigned long iterations_pi=MAXITER, prec_t maxresidual_pi=SOLPREC,
                unsigned long iterations_vi=MAXITER, prec_t maxresidual_vi=SOLPREC/2,
                bool print_progress=false) {
    assert(nature.size() == thresholds.size());
    assert(nature.size() == mdp.state_count());


    return mpi_jac<SType, PolicyNature<T>>(mdp, discount, valuefunction, 
                    PolicyNature<T>(policy,internal::zip(nature,thresholds)), 
                    iterations_pi, maxresidual_pi,
                    iterations_vi, maxresidual_vi, 
                    print_progress);
}

template<class SType, class T = prec_t>
inline auto rsolve_mpi(const GRMDP<SType>& mdp, prec_t discount,
                const NatureResponse<T>& nature, const vector<T>& thresholds,
                const numvec& valuefunction=numvec(0), const indvec& policy = indvec(0),
                unsigned long iterations_pi=MAXITER, prec_t maxresidual_pi=SOLPREC,
                unsigned long iterations_vi=MAXITER, prec_t maxresidual_vi=SOLPREC/2,
                bool print_progress=false) {
    assert(nature.size() == thresholds.size());
    assert(nature.size() == mdp.state_count());


    return mpi_jac<SType, PolicyNature<T>>(mdp, discount, valuefunction, 
                    PolicyNature<T>(policy,internal::zip(nature,thresholds)), 
                    iterations_pi, maxresidual_pi,
                    iterations_vi, maxresidual_vi, 
                    print_progress);
}

inline NatureResponse<prec_t> string_to_nature(string nature){
    if(nature == "robust_unbounded") return robust_unbounded;
    if(nature == "optimistic_unbounded") return optimistic_unbounded;
    if(nature == "robust_l1") return robust_l1;
    if(nature == "optimistic_l1") return optimistic_l1;
    throw invalid_argument("Unknown nature.");
}



}}