memory unfolder can label states with the memory state

src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp

@@ -11,7 +11,8 @@ namespace storm {
 
             
             template<typename ValueType>
-            PomdpMemoryUnfolder<ValueType>::PomdpMemoryUnfolder(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory) : pomdp(pomdp), memory(memory) {
+            PomdpMemoryUnfolder<ValueType>::PomdpMemoryUnfolder(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory, bool addMemoryLabels)
+            : pomdp(pomdp), memory(memory), addMemoryLabels(addMemoryLabels) {
                 // intentionally left empty
             }
 
@@ -94,6 +95,15 @@ namespace storm {
                     }
                     labeling.addLabel(labelName, std::move(newStates));
                 }
+                if (addMemoryLabels) {
+                    for (uint64_t memState = 0; memState < memory.getNumberOfStates(); ++memState) {
+                        storm::storage::BitVector newStates(pomdp.getNumberOfStates() * memory.getNumberOfStates(), false);
+                        for (uint64_t modelState = 0; modelState < pomdp.getNumberOfStates(); ++modelState) {
+                            newStates.set(getUnfoldingState(modelState, memState));
+                        }
+                        labeling.addLabel("memstate_"+std::to_string(memState), newStates);
+                    }
+                }
                 return labeling;
             }
         

src/storm-pomdp/transformer/PomdpMemoryUnfolder.h

@@ -12,7 +12,7 @@ namespace storm {
 
         public:
             
-            PomdpMemoryUnfolder(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory);
+            PomdpMemoryUnfolder(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory, bool addMemoryLabels = false);
             
             std::shared_ptr<storm::models::sparse::Pomdp<ValueType>> transform() const;
 
@@ -33,6 +33,8 @@ namespace storm {
             
             storm::models::sparse::Pomdp<ValueType> const& pomdp;
             storm::storage::PomdpMemory const& memory;
+
+            bool addMemoryLabels;
         };
     }
 }