diff --git a/src/storage/model.cpp b/src/storage/model.cpp index 167e003..e2c82fd 100644 --- a/src/storage/model.cpp +++ b/src/storage/model.cpp @@ -156,6 +156,7 @@ void define_model(py::module& m) { .def_property_readonly("has_transition_rewards", &RewardModel::hasTransitionRewards) .def_property_readonly("transition_rewards", [](RewardModel& rewardModel) {return rewardModel.getTransitionRewardMatrix();}) .def_property_readonly("state_rewards", [](RewardModel& rewardModel) {return rewardModel.getStateRewardVector();}) + .def("get_state_reward", [](RewardModel& rewardModel, uint64_t state) {return rewardModel.getStateReward(state);}) .def_property_readonly("state_action_rewards", [](RewardModel& rewardModel) {return rewardModel.getStateActionRewardVector();}) .def("reduce_to_state_based_rewards", [](RewardModel& rewardModel, SparseMatrix const& transitions, bool onlyStateRewards){return rewardModel.reduceToStateBasedRewards(transitions, onlyStateRewards);}, py::arg("transition_matrix"), py::arg("only_state_rewards"), "Reduce to state-based rewards") ; @@ -201,6 +202,7 @@ void define_model(py::module& m) { .def_property_readonly("has_transition_rewards", &RewardModel::hasTransitionRewards) .def_property_readonly("transition_rewards", [](RewardModel& rewardModel) {return rewardModel.getTransitionRewardMatrix();}) .def_property_readonly("state_rewards", [](RewardModel& rewardModel) {return rewardModel.getStateRewardVector();}) + .def("get_state_reward", [](RewardModel& rewardModel, uint64_t state) {return rewardModel.getStateReward(state);}) .def_property_readonly("state_action_rewards", [](RewardModel& rewardModel) {return rewardModel.getStateActionRewardVector();}) .def("reduce_to_state_based_rewards", [](RewardModel& rewardModel, SparseMatrix const& transitions, bool onlyStateRewards){return rewardModel.reduceToStateBasedRewards(transitions, onlyStateRewards);}, py::arg("transition_matrix"), py::arg("only_state_rewards"), "Reduce to state-based rewards") ;