Simulation.hpp

#pragma once

#include "Samples.hpp"
#include "definitions.hpp"

#include <utility>
#include <vector>
#include <memory>
#include <random>
#include <functional>
#include <cmath>
#include <algorithm>
#include <cmath>
#include <string>

#include "cpp11-range-master/range.hpp"


namespace craam{
namespace msen {

using namespace std;
using namespace util::lang;











template<class Sim, class SampleType=Samples<typename Sim::State, typename Sim::Action>>
void simulate(
            Sim& sim, SampleType& samples,
            const function<typename Sim::Action(typename Sim::State&)>& policy,
            long horizon, long runs, long tran_limit=-1, prec_t prob_term=0.0,
            random_device::result_type seed = random_device{}()){

    long transitions = 0;

    // initialize random numbers to be used with random termination
    default_random_engine generator(seed);
    uniform_real_distribution<double> distribution(0.0,1.0);

    for(auto run=0l; run < runs; run++){

        typename Sim::State state = sim.init_state();
        samples.add_initial(state);

        for(auto step : range(0l,horizon)){
            // check form termination conditions
            if(sim.end_condition(state) || (tran_limit > 0 && transitions > tran_limit) )
                break;

            auto action = policy(state);
            auto reward_state = sim.transition(state,action);

            auto reward = reward_state.first;
            auto nextstate = move(reward_state.second);

            samples.add_sample(move(state), move(action), nextstate, reward, 1.0, step, run);
            state = move(nextstate);

            // test the termination probability only after at least one transition
            if( (prob_term > 0.0) && (distribution(generator) <= prob_term) )
                break;
            transitions++;
        };

        if(tran_limit > 0 && transitions > tran_limit)
            break;
    }
}

template<class Sim, class SampleType=Samples<typename Sim::State, typename Sim::Action>>
SampleType simulate(
            Sim& sim,
            const function<typename Sim::Action(typename Sim::State&)>& policy,
            long horizon, long runs, long tran_limit=-1, prec_t prob_term=0.0,
            random_device::result_type seed = random_device{}()){

    SampleType samples = SampleType();
    simulate(sim, samples, policy, horizon, runs, tran_limit, prob_term, seed);
    return samples;
}


template<class Sim>
pair<vector<typename Sim::State>, numvec> 
simulate_return(Sim& sim, prec_t discount,
                const function<typename Sim::Action(typename Sim::State&)>& policy,
                long horizon, long runs, prec_t prob_term=0.0,
                random_device::result_type seed = random_device{}()){

    long transitions = 0;
    // initialize random numbers to be used with random termination
    default_random_engine generator(seed);
    uniform_real_distribution<double> distribution(0.0,1.0);

    // pre-initialize output values
    vector<typename Sim::State> start_states(runs);
    numvec returns(runs);

    for(auto run : range(0l,runs)){
        typename Sim::State state = sim.init_state();
        start_states[run] = state;

        prec_t runreturn = 0;
        for(auto step : range(0l,horizon)){
            // check from-state termination conditions
            if(sim.end_condition(state))
                break;

            auto action = policy(state);
            auto reward_state = sim.transition(state,action);
            
            auto reward = reward_state.first;
            auto nextstate = move(reward_state.second);

            runreturn += reward * pow(discount, step);
            state = move(nextstate);
            // test the termination probability only after at least one transition
            if( (prob_term > 0.0) && (distribution(generator) <= prob_term) )
                break;
            transitions++;
        };
        returns[run] = runreturn;
    }
    return make_pair(move(start_states), move(returns));
}

// ************************************************************************************
// **** Random(ized) policies ****
// ************************************************************************************

template<class Sim>
class RandomPolicy{

public:
    using State = typename Sim::State;
    using Action = typename Sim::Action;

    RandomPolicy(const Sim& sim, random_device::result_type seed = random_device{}()) : 
                sim(sim), gen(seed){};

    Action operator() (State state){
        uniform_int_distribution<long> dst(0,sim.action_count(state)-1);
        return sim.action(state,dst(gen));
    };

private:
    const Sim& sim;
    default_random_engine gen;
};

template<typename Sim>
class RandomizedPolicy{

public:
    using State = typename Sim::State;
    using Action = typename Sim::Action;

    RandomizedPolicy(const Sim& sim, const vector<numvec>& probabilities,random_device::result_type seed = random_device{}()):
        gen(seed), distributions(probabilities.size()), sim(sim){

        for(auto pi : indices(probabilities)){
            
            // check that this distribution is correct
            const numvec& prob = probabilities[pi];
            prec_t sum = accumulate(prob.begin(), prob.end(), 0.0);
            
            if(abs(sum - 1) > SOLPREC){
                throw invalid_argument("Action probabilities must sum to 1 in state " + to_string(pi));
            } 
            distributions[pi] = discrete_distribution<long>(prob.begin(), prob.end());
        }
    };

    Action operator() (State state){
        // check that the state is valid for this policy
        long sl = static_cast<long>(state);
        assert(sl >= 0 && size_t(sl) < distributions.size());

        auto& dst = distributions[sl];
        // existence of the action is check by the simulator
        return sim.action(state,dst(gen));
    };

protected:

    default_random_engine gen;

    vector<discrete_distribution<long>> distributions;

    const Sim& sim;
};


template<typename Sim>
class DeterministicPolicy{

public:
    using State = typename Sim::State;
    using Action = typename Sim::Action;

    DeterministicPolicy(const Sim& sim, indvec actions):
        actions(actions), sim(sim) {};

    Action operator() (State state){
        // check that the state is valid for this policy
        long sl = static_cast<long>(state);

        assert(sl >= 0 && size_t(sl) < actions.size());

        // existence of the action is checked by the simulator
        return sim.action(state,actions[sl]);
    };

protected:
    indvec actions;

    const Sim& sim;
};



// ************************************************************************************
// **** MDP simulation ****
// ************************************************************************************

class ModelSimulator{

public:
    typedef long State;
    typedef long Action;

    ModelSimulator(const shared_ptr<const MDP>& mdp, const Transition& initial, 
                        random_device::result_type seed = random_device{}()) :
                gen(seed), mdp(mdp), initial(initial){

        if(abs(initial.sum_probabilities() - 1) > SOLPREC)
            throw invalid_argument("Initial transition probabilities must sum to 1");
    }

    ModelSimulator(const shared_ptr<MDP>& mdp, const Transition& initial,random_device::result_type seed = random_device{}()) : 
        ModelSimulator(const_pointer_cast<const MDP>(mdp), initial, seed) {};

    State init_state(){
        const numvec& probs = initial.get_probabilities();
        const indvec& inds = initial.get_indices();
        auto dst = discrete_distribution<long>(probs.begin(), probs.end());
        return inds[dst(gen)];
    }
    
    pair<double,State> transition(State state, Action action){

        assert(state >= 0 && size_t(state) < mdp->size());
        const auto& mdpstate = (*mdp)[state];

        assert(action >= 0 && size_t(action) < mdpstate.size());
        const auto& mdpaction = mdpstate[action];

        if(!mdpstate.is_valid(action))
            throw invalid_argument("Cannot transition using an invalid action");

        const auto& tran = mdpaction.get_outcome();

        const numvec& probs = tran.get_probabilities();
        const numvec& rews = tran.get_rewards();
        const indvec& inds = tran.get_indices();

        // check if the transition sums to 1, if not use the remainder 
        // as a probability of terminating
        prec_t prob_termination = 1 - tran.sum_probabilities();

        discrete_distribution<long> dst;

        if(prob_termination > SOLPREC){
            // copy the probabilities (there should be a faster way too)
            numvec copy_probs(probs);
            copy_probs.push_back(prob_termination);

            dst = discrete_distribution<long>(copy_probs.begin(), copy_probs.end());
        }else{
            dst = discrete_distribution<long>(probs.begin(), probs.end());
        }

        const size_t nextindex = dst(gen);

        // check if need to transition to a terminal state
        const State nextstate = nextindex < inds.size() ? 
                                inds[nextindex] : mdp->size();

        // reward is zero when transitioning to a terminal state
        const prec_t reward = nextindex < inds.size() ? 
                                rews[nextindex] : 0.0;

        return make_pair(reward, nextstate);
    }

    bool end_condition(State s) const 
        {return (size_t(s) >= mdp->size()) || (action_count(s) == 0);};

    size_t action_count(State state) const 
        {return (*mdp)[state].size();};
    
    Action action(State, long index) const
        {return index;};

protected:
    default_random_engine gen;

    shared_ptr<const MDP> mdp;
    
    Transition initial;
};

using ModelRandomPolicy = RandomPolicy<ModelSimulator>;

using ModelRandomizedPolicy = RandomizedPolicy<ModelSimulator>;

using ModelDeterministicPolicy = DeterministicPolicy<ModelSimulator>;


} // end namespace msen
} // end namespace craam