From d2cd142dfba86735fa0b15434f30c76d9c04db75 Mon Sep 17 00:00:00 2001
From: TimQu <tim.quatmann@cs.rwth-aachen.de>
Date: Tue, 27 Nov 2018 20:36:37 +0100
Subject: [PATCH] Fixed initial partitioning in sparse bisimulation with
 action-based rewards.

---
 .../BisimulationDecomposition.cpp             | 25 ++++++++++++++++---
 .../bisimulation/BisimulationDecomposition.h  |  5 ++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/storm/storage/bisimulation/BisimulationDecomposition.cpp b/src/storm/storage/bisimulation/BisimulationDecomposition.cpp
index 0b2578626..577e7e487 100644
--- a/src/storm/storage/bisimulation/BisimulationDecomposition.cpp
+++ b/src/storm/storage/bisimulation/BisimulationDecomposition.cpp
@@ -168,7 +168,7 @@ namespace storm {
         template<typename ModelType, typename BlockDataType>
         BisimulationDecomposition<ModelType, BlockDataType>::BisimulationDecomposition(ModelType const& model, storm::storage::SparseMatrix<ValueType> const& backwardTransitions, Options const& options) : model(model), backwardTransitions(backwardTransitions), options(options), partition(), comparator(), quotient(nullptr) {
             STORM_LOG_THROW(!options.getKeepRewards() || !model.hasRewardModel() || model.hasUniqueRewardModel(), storm::exceptions::IllegalFunctionCallException, "Bisimulation currently only supports models with at most one reward model.");
-            STORM_LOG_THROW(!options.getKeepRewards() || !model.hasRewardModel() || !model.getUniqueRewardModel().hasTransitionRewards(), storm::exceptions::IllegalFunctionCallException, "Bisimulation is currently supported for models with state or action rewards rewards only. Consider converting the transition rewards to state rewards (via suitable function calls).");
+            STORM_LOG_THROW(!options.getKeepRewards() || !model.hasRewardModel() || !model.getUniqueRewardModel().hasTransitionRewards(), storm::exceptions::IllegalFunctionCallException, "Bisimulation is currently supported for models with state or action rewards only. Consider converting the transition rewards to state rewards (via suitable function calls).");
             STORM_LOG_THROW(options.getType() != BisimulationType::Weak || !options.getBounded(), storm::exceptions::IllegalFunctionCallException, "Weak bisimulation cannot preserve bounded properties.");
             
             // Fix the respected atomic propositions if they were not explicitly given.
@@ -190,6 +190,7 @@ namespace storm {
             } else {
                 this->initializeLabelBasedPartition();
             }
+            STORM_LOG_WARN_COND(partition.size() > 1, "Initial partition consists only of a single block.");
             std::chrono::high_resolution_clock::duration initialPartitionTime = std::chrono::high_resolution_clock::now() - initialPartitionStart;
             
             this->initialize();
@@ -264,8 +265,21 @@ namespace storm {
             if (rewardModel.hasStateRewards()) {
                 this->splitInitialPartitionBasedOnRewards(rewardModel.getStateRewardVector());
             }
-            if (rewardModel.hasStateActionRewards() && (model.isOfType(storm::models::ModelType::Dtmc) || model.isOfType(storm::models::ModelType::Ctmc))) {
-                this->splitInitialPartitionBasedOnRewards(rewardModel.getStateActionRewardVector());
+            if (rewardModel.hasStateActionRewards()) {
+                if (model.isNondeterministicModel()) {
+                    std::vector<std::set<ValueType>> actionRewards;
+                    actionRewards.reserve(model.getNumberOfStates());
+                    for (storm::storage::sparse::state_type state = 0; state <  model.getNumberOfStates(); ++state) {
+                        std::set<ValueType> rewardsAtState;
+                        for (auto choice = model.getTransitionMatrix().getRowGroupIndices()[state]; choice < model.getTransitionMatrix().getRowGroupIndices()[state + 1]; ++choice) {
+                            rewardsAtState.insert(rewardModel.getStateActionReward(choice));
+                        }
+                        actionRewards.push_back(std::move(rewardsAtState));
+                    }
+                    this->splitInitialPartitionBasedOnActionRewards(actionRewards);
+                } else {
+                    this->splitInitialPartitionBasedOnRewards(rewardModel.getStateActionRewardVector());
+                }
             }
         }
         
@@ -274,6 +288,11 @@ namespace storm {
             partition.split([&rewardVector] (storm::storage::sparse::state_type const& a, storm::storage::sparse::state_type const& b) { return rewardVector[a] < rewardVector[b]; });
         }
         
+        template<typename ModelType, typename BlockDataType>
+        void BisimulationDecomposition<ModelType, BlockDataType>::splitInitialPartitionBasedOnActionRewards(std::vector<std::set<ValueType>> const& actionRewards) {
+            partition.split([&actionRewards] (storm::storage::sparse::state_type const& a, storm::storage::sparse::state_type const& b) { return actionRewards[a] < actionRewards[b]; });
+        }
+        
         template<typename ModelType, typename BlockDataType>
         void BisimulationDecomposition<ModelType, BlockDataType>::initializeLabelBasedPartition() {
             partition = storm::storage::bisimulation::Partition<BlockDataType>(model.getNumberOfStates());
diff --git a/src/storm/storage/bisimulation/BisimulationDecomposition.h b/src/storm/storage/bisimulation/BisimulationDecomposition.h
index dabdaf549..7e8d871fa 100644
--- a/src/storm/storage/bisimulation/BisimulationDecomposition.h
+++ b/src/storm/storage/bisimulation/BisimulationDecomposition.h
@@ -260,6 +260,11 @@ namespace storm {
              * Splits the initial partition based on the given reward vector.
              */
             virtual void splitInitialPartitionBasedOnRewards(std::vector<ValueType> const& rewardVector);
+
+            /*!
+             * Splits the initial partition based on the given vector of action rewards.
+             */
+            virtual void splitInitialPartitionBasedOnActionRewards(std::vector<std::set<ValueType>> const& rewardVector);
             
             /*!
              * Constructs the blocks of the decomposition object based on the current partition.