diff --git a/src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp b/src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp index 086a1e389..6687c30db 100644 --- a/src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp +++ b/src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp @@ -11,7 +11,8 @@ namespace storm { template - PomdpMemoryUnfolder::PomdpMemoryUnfolder(storm::models::sparse::Pomdp const& pomdp, storm::storage::PomdpMemory const& memory) : pomdp(pomdp), memory(memory) { + PomdpMemoryUnfolder::PomdpMemoryUnfolder(storm::models::sparse::Pomdp 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; } diff --git a/src/storm-pomdp/transformer/PomdpMemoryUnfolder.h b/src/storm-pomdp/transformer/PomdpMemoryUnfolder.h index 69b18559a..41f8d0526 100644 --- a/src/storm-pomdp/transformer/PomdpMemoryUnfolder.h +++ b/src/storm-pomdp/transformer/PomdpMemoryUnfolder.h @@ -12,7 +12,7 @@ namespace storm { public: - PomdpMemoryUnfolder(storm::models::sparse::Pomdp const& pomdp, storm::storage::PomdpMemory const& memory); + PomdpMemoryUnfolder(storm::models::sparse::Pomdp const& pomdp, storm::storage::PomdpMemory const& memory, bool addMemoryLabels = false); std::shared_ptr> transform() const; @@ -33,6 +33,8 @@ namespace storm { storm::models::sparse::Pomdp const& pomdp; storm::storage::PomdpMemory const& memory; + + bool addMemoryLabels; }; } } \ No newline at end of file