diff --git a/src/storage/model.cpp b/src/storage/model.cpp index e2c82fd..ec3d43f 100644 --- a/src/storage/model.cpp +++ b/src/storage/model.cpp @@ -157,7 +157,8 @@ void define_model(py::module& m) { .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("get_state_action_reward", [](RewardModel& rewardModel, uint64_t action_index) {return rewardModel.getStateActionReward(action_index);}) + .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") ; @@ -203,6 +204,8 @@ void define_model(py::module& m) { .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("get_state_action_reward", [](RewardModel& rewardModel, uint64_t action_index) {return rewardModel.getStateActionReward(action_index);}) + .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") ;