From cf3bfc3d2ddc90c80735f030be69d65e0d496887 Mon Sep 17 00:00:00 2001
From: Sebastian Junges <sebastian.junges@gmail.com>
Date: Thu, 18 Jun 2020 12:26:51 -0700
Subject: [PATCH] simulator presents rewards for efficiency

---
 .../DiscreteTimeSparseModelSimulator.cpp      | 40 +++++++++++++++++--
 .../DiscreteTimeSparseModelSimulator.h        |  5 ++-
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/src/storm/simulator/DiscreteTimeSparseModelSimulator.cpp b/src/storm/simulator/DiscreteTimeSparseModelSimulator.cpp
index 2e9161b8a..d6669516d 100644
--- a/src/storm/simulator/DiscreteTimeSparseModelSimulator.cpp
+++ b/src/storm/simulator/DiscreteTimeSparseModelSimulator.cpp
@@ -4,9 +4,16 @@
 namespace storm {
     namespace simulator {
         template<typename ValueType, typename RewardModelType>
-        DiscreteTimeSparseModelSimulator<ValueType,RewardModelType>::DiscreteTimeSparseModelSimulator(storm::models::sparse::Model<ValueType, RewardModelType> const& model) : currentState(*model.getInitialStates().begin()), model(model) {
+        DiscreteTimeSparseModelSimulator<ValueType,RewardModelType>::DiscreteTimeSparseModelSimulator(storm::models::sparse::Model<ValueType, RewardModelType> const& model) : model(model), currentState(*model.getInitialStates().begin()), zeroRewards(model.getNumberOfRewardModels(), storm::utility::zero<ValueType>()) {
             STORM_LOG_WARN_COND(model.getInitialStates().getNumberOfSetBits()==1, "The model has multiple initial states. This simulator assumes it starts from the initial state with the lowest index.");
-
+            lastRewards = zeroRewards;
+            uint64_t i = 0;
+            for (auto const& rewModPair : model.getRewardModels()) {
+                if (rewModPair.second.hasStateRewards()) {
+                    lastRewards[i] += rewModPair.second.getStateReward(currentState);
+                }
+                ++i;
+            }
         }
 
         template<typename ValueType, typename RewardModelType>
@@ -18,19 +25,34 @@ namespace storm {
         bool DiscreteTimeSparseModelSimulator<ValueType,RewardModelType>::step(uint64_t action) {
             // TODO lots of optimization potential.
             //  E.g., do not sample random numbers if there is only a single transition
+            lastRewards = zeroRewards;
             ValueType probability = generator.random();
             STORM_LOG_ASSERT(action < model.getTransitionMatrix().getRowGroupSize(currentState), "Action index higher than number of actions");
             uint64_t row = model.getTransitionMatrix().getRowGroupIndices()[currentState] + action;
+            uint64_t i = 0;
+            for (auto const& rewModPair : model.getRewardModels()) {
+                if (rewModPair.second.hasStateActionRewards()) {
+                    lastRewards[i] += rewModPair.second.getStateActionReward(row);
+                }
+                ++i;
+            }
             ValueType sum = storm::utility::zero<ValueType>();
             for (auto const& entry : model.getTransitionMatrix().getRow(row)) {
                 sum += entry.getValue();
                 if (sum >= probability) {
                     currentState = entry.getColumn();
+                    i = 0;
+                    for (auto const& rewModPair : model.getRewardModels()) {
+                        if (rewModPair.second.hasStateRewards()) {
+                            lastRewards[i] += rewModPair.second.getStateReward(currentState);
+                        }
+                        ++i;
+                    }
                     return true;
                 }
             }
+            // This position should never be reached
             return false;
-            STORM_LOG_ASSERT(false, "This position should never be reached");
         }
 
         template<typename ValueType, typename RewardModelType>
@@ -41,9 +63,21 @@ namespace storm {
         template<typename ValueType, typename RewardModelType>
         bool DiscreteTimeSparseModelSimulator<ValueType,RewardModelType>::resetToInitial() {
             currentState = *model.getInitialStates().begin();
+            lastRewards = zeroRewards;
+            uint64_t i = 0;
+            for (auto const& rewModPair : model.getRewardModels()) {
+                if (rewModPair.second.hasStateRewards()) {
+                    lastRewards[i] += rewModPair.second.getStateReward(currentState);
+                }
+                ++i;
+            }
             return true;
         }
 
+        template<typename ValueType, typename RewardModelType>
+        std::vector<ValueType> const& DiscreteTimeSparseModelSimulator<ValueType,RewardModelType>::getLastRewards() const {
+            return lastRewards;
+        }
 
         template class DiscreteTimeSparseModelSimulator<double>;
     }
diff --git a/src/storm/simulator/DiscreteTimeSparseModelSimulator.h b/src/storm/simulator/DiscreteTimeSparseModelSimulator.h
index 1637bffec..a0dcf4e3a 100644
--- a/src/storm/simulator/DiscreteTimeSparseModelSimulator.h
+++ b/src/storm/simulator/DiscreteTimeSparseModelSimulator.h
@@ -19,11 +19,14 @@ namespace storm {
             DiscreteTimeSparseModelSimulator(storm::models::sparse::Model<ValueType, RewardModelType> const& model);
             void setSeed(uint64_t);
             bool step(uint64_t action);
+            std::vector<ValueType> const& getLastRewards() const;
             uint64_t getCurrentState() const;
             bool resetToInitial();
         protected:
-            uint64_t currentState;
             storm::models::sparse::Model<ValueType, RewardModelType> const& model;
+            uint64_t currentState;
+            std::vector<ValueType> lastRewards;
+            std::vector<ValueType> zeroRewards;
             storm::utility::RandomProbabilityGenerator<ValueType> generator;
         };
     }