occupancies.hpp

#pragma once

#include "../RMDP.hpp"

#include <Eigen/Dense>
#include "../cpp11-range-master/range.hpp"

namespace craam{namespace algorithms{

using namespace std;
using namespace Eigen;

namespace internal{

    template<class SType>
    inline Transition mean_transition_state(const SType& state, long index, const pair<indvec,vector<numvec>>& policies){
        return state.mean_transition(policies.first[index], policies.second[index]);
    }

    template<class SType>
    inline Transition mean_transition_state(const SType& state, long index, const indvec& policy){
        return state.mean_transition(policy[index]);
    }

    template<class SType>
    inline prec_t mean_reward_state(const SType& state, long index, const pair<indvec,vector<numvec>>& policies){
        return state.mean_reward(policies.first[index], policies.second[index]);
    }

    template<class SType>
    inline prec_t mean_reward_state(const SType& state, long index, const indvec& policy){
        return state.mean_reward(policy[index]);
    }
}

template<typename SType, typename Policies> 
inline MatrixXd transition_mat(const GRMDP<SType>& rmdp, const Policies& policies, bool transpose = false) {
    const size_t n = rmdp.state_count();
    MatrixXd result = MatrixXd::Zero(n,n);

    const auto& states = rmdp.get_states();
    #pragma omp parallel for
    for(size_t s = 0; s < n; s++){
        const Transition&& t = internal::mean_transition_state(states[s], s, policies);

        const auto& indexes = t.get_indices();
        const auto& probabilities = t.get_probabilities();

        if(!transpose){
            for(size_t j=0; j < t.size(); j++)
                result(s,indexes[j]) = probabilities[j];
        }else{
            for(size_t j=0; j < t.size(); j++)
                result(indexes[j],s) = probabilities[j];
        }
    }
    return result;
}

template<typename SType, typename Policy>
inline numvec rewards_vec(const GRMDP<SType>& rmdp, const Policy& policies){
    
    const auto n = rmdp.state_count();
    numvec rewards(n);

    #pragma omp parallel for
    for(size_t s=0; s < n; s++){
        const SType& state = rmdp[s];
        if(state.is_terminal())
            rewards[s] = 0;
        else
            rewards[s] = internal::mean_reward_state(state, s, policies);
    }
    return rewards;
}

template<typename SType, typename Policies>
inline numvec 
occfreq_mat(const GRMDP<SType>& rmdp, const Transition& init, prec_t discount,
                 const Policies& policies) {
    const auto n = rmdp.state_count();

    // initial distribution
    const numvec& ivec = init.probabilities_vector(n);
    const VectorXd initial_vec = Map<const VectorXd,Unaligned>(ivec.data(),ivec.size());

    // get transition matrix and construct (I - gamma * P^T)
    MatrixXd t_mat = MatrixXd::Identity(n,n)  - discount * transition_mat(rmdp, policies, true);

    // solve set of linear equations
    numvec result(n,0);
    Map<VectorXd,Unaligned>(result.data(),result.size()) = HouseholderQR<MatrixXd>(t_mat).solve(initial_vec);

    return result;
}

}}