diff --git a/src/pomdp/tracker.cpp b/src/pomdp/tracker.cpp index 766d03c..63c9d33 100644 --- a/src/pomdp/tracker.cpp +++ b/src/pomdp/tracker.cpp @@ -1,15 +1,33 @@ #include "tracker.h" #include "src/helpers.h" #include +#include template using SparsePomdp = storm::models::sparse::Pomdp; -template using SparsePomdpTracker = storm::generator::BeliefSupportTracker; +template using SparsePomdpTracker = storm::generator::BeliefSupportTracker; +template using NDPomdpTrackerSparse = storm::generator::NondeterministicBeliefTracker>; void define_tracker(py::module& m) { py::class_> tracker(m, "BeliefSupportTrackerDouble", "Tracker for BeliefSupports"); tracker.def(py::init const&>(), py::arg("pomdp")); tracker.def("get_current_belief_support", &SparsePomdpTracker::getCurrentBeliefSupport, "What is the support given the trace so far"); tracker.def("track", &SparsePomdpTracker::track, py::arg("action"), py::arg("observation")); + + py::class_> sbel(m, "SparseBeliefStateDouble", "Belief state in sparse format"); + sbel.def("get", &storm::generator::SparseBeliefState::get, py::arg("state")); + sbel.def_property_readonly("risk", &storm::generator::SparseBeliefState::getRisk); + sbel.def("__str__", &storm::generator::SparseBeliefState::toString); + sbel.def_property_readonly("is_valid", &storm::generator::SparseBeliefState::isValid); + + py::class_> ndetbelieftracker(m, "NondeterministicBeliefTrackerDoubleSparse", "Tracker for belief states and uncontrollable actions"); + ndetbelieftracker.def(py::init const&>(), py::arg("pomdp")); + ndetbelieftracker.def("reset", &NDPomdpTrackerSparse::reset); + ndetbelieftracker.def("set_risk", &NDPomdpTrackerSparse::setRisk, py::arg("risk")); + ndetbelieftracker.def("obtain_current_risk",&NDPomdpTrackerSparse::getCurrentRisk, py::arg("max")=true); + ndetbelieftracker.def("track", &NDPomdpTrackerSparse::track, py::arg("observation")); + ndetbelieftracker.def("obtain_beliefs", &NDPomdpTrackerSparse::getCurrentBeliefs); + ndetbelieftracker.def("obtain_last_observation", &NDPomdpTrackerSparse::getCurrentObservation); + } \ No newline at end of file