From f0974c8f4e49732c644763619a4d0073617ac88b Mon Sep 17 00:00:00 2001
From: YouGuessedMyName <ivo@melse.info>
Date: Mon, 9 Dec 2024 15:35:40 +0100
Subject: [PATCH] progress removing actions

---
 stormvogel/mapping.py        |   6 +-
 stormvogel/model.py          | 112 ++++++++++----------
 stormvogel/result.py         |   2 +-
 stormvogel/simulator.py      |   9 +-
 stormvogel/visualization.py  |   6 +-
 tests/saved_test_layout.json |   2 +-
 tests/test_mapping.py        |  78 +++++++-------
 tests/test_model_methods.py  |  63 +++++------
 tests/test_result.py         |   2 +-
 tests/test_simulator.py      | 196 ++++++++++++++++++-----------------
 tests/test_visualization.py  |  48 ++++-----
 11 files changed, 264 insertions(+), 260 deletions(-)

diff --git a/stormvogel/mapping.py b/stormvogel/mapping.py
index 2bc9e46..c1c4322 100644
--- a/stormvogel/mapping.py
+++ b/stormvogel/mapping.py
@@ -413,7 +413,7 @@ def map_mdp(sparsemdp: stormpy.storage.SparseDtmc) -> stormvogel.model.Model:
                 else:
                     actionlabels = frozenset()
                 # TODO assign the correct action name and not only an index
-                action = model.new_action(str(i), actionlabels)
+                action = model.new_action(actionlabels)
                 branch = [(x.value(), model.get_state_by_id(x.column)) for x in row]
                 transition[action] = stormvogel.model.Branch(branch)
                 transitions = stormvogel.model.Transition(transition)
@@ -492,7 +492,7 @@ def map_pomdp(sparsepomdp: stormpy.storage.SparsePomdp) -> stormvogel.model.Mode
                     actionlabels = frozenset()
 
                 # TODO assign the correct action name and not only an index
-                action = model.new_action(str(i), actionlabels)
+                action = model.new_action(actionlabels)
                 branch = [(x.value(), model.get_state_by_id(x.column)) for x in row]
                 transition[action] = stormvogel.model.Branch(branch)
                 transitions = stormvogel.model.Transition(transition)
@@ -540,7 +540,7 @@ def map_ma(sparsema: stormpy.storage.SparseMA) -> stormvogel.model.Model:
                     actionlabels = frozenset()
 
                 # TODO assign the correct action name and not only an index
-                action = model.new_action(str(i), actionlabels)
+                action = model.new_action(actionlabels)
                 branch = [(x.value(), model.get_state_by_id(x.column)) for x in row]
                 transition[action] = stormvogel.model.Branch(branch)
                 transitions = stormvogel.model.Transition(transition)
diff --git a/stormvogel/model.py b/stormvogel/model.py
index e2bec1f..6be2e10 100644
--- a/stormvogel/model.py
+++ b/stormvogel/model.py
@@ -199,36 +199,29 @@ class Action:
     """Represents an action, e.g., in MDPs.
         Note that this action object is completely independent of its corresponding branch.
         Their relation is managed by Transitions.
+        Two actions with the same labels are considered equal.
 
     Args:
         name: A name for this action.
         labels: The labels of this action. Corresponds to Storm labels.
     """
+    @staticmethod 
+    def create(labels: frozenset[str] | str | None = None) -> 'Action':
+        if isinstance(labels, str):
+            return Action(frozenset({labels}))
+        elif isinstance(labels, frozenset):
+            return Action(labels)
+        else:
+            return Action(frozenset())
 
-    name: str
     labels: frozenset[str]
 
     def __str__(self):
-        return f"Action {self.name} with labels {self.labels}"
-
-    # TODO remove these after modifying the whole code base to remove names.
-    def __eq__(self, other):
-        if isinstance(other, Action):
-            return self.labels == other.labels
-        return False
-
-    def strict_eq(self, other):
-        """Also requires the names to be equal."""
-        if isinstance(other, Action):
-            return self.name == other.name and self.labels == other.labels
-        return False
-
-    # def __hash__(self):
-    #     return self.labels.__hash__()
+        return f"Action with labels {self.labels}"
 
 
 # The empty action. Used for DTMCs and empty action transitions in mdps.
-EmptyAction = Action("empty", frozenset())
+EmptyAction = Action(frozenset())
 
 
 @dataclass(order=True)
@@ -371,21 +364,24 @@ def set_from_rewards_vector(self, vector: list[Number]) -> None:
                 self.rewards[s.id, a] = vector[combined_id]
                 combined_id += 1
 
-    def get_state_reward(self, state: State) -> Number:
-        """Gets the reward at said state or state action pair"""
+    def get_state_reward(self, state: State) -> Number | None:
+        """Gets the reward at said state or state action pair. Return None if no reward is present."""
         if self.model.supports_actions():
             RuntimeError(
                 "This is a model with actions. Please call the get_action_state_reward(_at_id) function instead"
             )
-        return self.rewards[state.id, EmptyAction]
+        if (state.id, EmptyAction) in self.rewards:
+            return self.rewards[state.id, EmptyAction]
+        else:
+            return None
 
     def get_state_action_reward(self, state: State, action: Action) -> Number | None:
         """Gets the reward at said state or state action pair. Returns None if no reward was found."""
         if self.model.supports_actions():
             if action in state.available_actions():
-                try:
+                if (state.id, action) in self.rewards:
                     return self.rewards[state.id, action]
-                except KeyError:
+                else:
                     return None
             else:
                 RuntimeError("This action is not available in this state")
@@ -474,7 +470,7 @@ class Model:
     # Both of these are hashed by the id of the state (=number in the matrix)
     states: dict[int, State]
     transitions: dict[int, Transition]
-    actions: dict[str, Action] | None
+    actions: set[Action] | None
     rewards: list[RewardModel]
     # In ctmcs we work with rate transitions but additionally we can optionally store exit rates (hashed by id of the state)
     exit_rates: dict[int, Number] | None
@@ -492,7 +488,7 @@ def __init__(
 
         # Initialize actions if those are supported by the model type
         if self.supports_actions():
-            self.actions = {}
+            self.actions = set()
         else:
             self.actions = None
 
@@ -625,7 +621,7 @@ def get_state_action_id(self, state: State, action: Action) -> int | None:
         for s in self.states.values():
             for a in s.available_actions():
                 if (
-                    a.name == action.name
+                    a == action
                     and action in s.available_actions()
                     and s == state
                 ):
@@ -668,6 +664,8 @@ def set_transitions(
         """Set the transition from a state."""
         if not isinstance(transitions, Transition):
             transitions = transition_from_shorthand(transitions)
+        if not self.actions is None and EmptyAction in transitions.transition.keys():
+            self.actions.add(EmptyAction)
         self.transitions[s.id] = transitions
 
     def add_transitions(
@@ -714,8 +712,8 @@ def add_transitions(
             else:
                 for action, branch in transitions.transition.items():
                     assert self.actions is not None
-                    if action not in self.actions.values():
-                        self.actions[action.name] = action
+                    if action not in self.actions:
+                        self.actions.add(action)
                     self.transitions[s.id].transition[action] = branch
 
     def get_transitions(self, state_or_id: State | int) -> Transition:
@@ -733,19 +731,15 @@ def get_branch(self, state_or_id: State | int) -> Branch:
             raise RuntimeError("Called get_branch on a non-empty transition.")
         return transition[EmptyAction]
 
-    def new_action(self, name: str, labels: frozenset[str] | None = None) -> Action:
+    def new_action(self, labels: frozenset[str] | str | None = None) -> Action:
         """Creates a new action and returns it."""
         if not self.supports_actions():
             raise RuntimeError(
                 "Called new_action on a model that does not support actions"
             )
         assert self.actions is not None
-        if name in self.actions:
-            raise RuntimeError(
-                f"Tried to add action {name} but that action already exists"
-            )
-        action = Action(name, labels if labels else frozenset())
-        self.actions[name] = action
+        action = Action.create(labels)
+        self.actions.add(action)
         return action
 
     def reassign_ids(self):
@@ -845,30 +839,32 @@ def remove_transitions_between_states(
                 "This method only works for models that don't support actions."
             )
 
-    def get_action(self, name: str) -> Action:
-        """Gets an existing action."""
-        if not self.supports_actions():
-            raise RuntimeError(
-                "Called get_action on a model that does not support actions"
-            )
-        assert self.actions is not None
-        if name not in self.actions:
-            raise RuntimeError(
-                f"Tried to get action {name} but that action does not exist"
-            )
-        return self.actions[name]
-
-    def action(self, name: str) -> Action:
+    # TODO possibly obsolete?
+    # def get_action(self, name: str) -> Action:
+    #     """Gets an existing action."""
+    #     if not self.supports_actions():
+    #         raise RuntimeError(
+    #             "Called get_action on a model that does not support actions"
+    #         )
+    #     assert self.actions is not None
+    #     if name not in self.actions:
+    #         raise RuntimeError(
+    #             f"Tried to get action {name} but that action does not exist"
+    #         )
+    #     return self.actions[name]
+
+    def action(self, labels: frozenset[str] | str | None) -> Action:
         """New action or get action if it exists."""
         if not self.supports_actions():
             raise RuntimeError(
-                "Called get_action on a model that does not support actions"
+                "Called method action on a model that does not support actions"
             )
         assert self.actions is not None
-        if name in self.actions:
-            return self.get_action(name)
-        else:
-            return self.new_action(name)
+        action = Action.create(labels)
+
+        if not action in self.actions:
+            self.new_action(labels)
+        return action
 
     def new_state(
         self,
@@ -980,7 +976,7 @@ def to_dot(self) -> str:
         for state_id, transition in self.transitions.items():
             for action, branch in transition.transition.items():
                 if action != EmptyAction:
-                    dot += f'{action.name.replace(" ", "_")}{state_id} [ label = "", shape=point ];\n'
+                    dot += f'{state_id} [ label = "", shape=point ];\n'
         for state_id, transition in self.transitions.items():
             for action, branch in transition.transition.items():
                 if action == EmptyAction:
@@ -989,9 +985,9 @@ def to_dot(self) -> str:
                         dot += f'{state_id} -> {target.id} [ label = "{prob}" ];\n'
                 else:
                     # Draw actions, then probabilities
-                    dot += f'{state_id} -> {action.name.replace(" ", "_")}{state_id} [ label = "{action.name}" ];\n'
+                    dot += f'{state_id} -> {state_id} [ label = "{action.name}" ];\n'
                     for prob, target in branch.branch:
-                        dot += f'{action.name.replace(" ", "_")}{state_id} -> {target.id} [ label = "{prob}" ];\n'
+                        dot += f'{state_id} -> {target.id} [ label = "{prob}" ];\n'
 
         dot += "}"
         return dot
@@ -1021,7 +1017,7 @@ def __eq__(self, other) -> bool:
             if self.supports_actions():
                 assert self.actions is not None and other.actions is not None
                 for action, other_action in zip(
-                    sorted(self.actions.values()), sorted(other.actions.values())
+                    sorted(self.actions), sorted(other.actions)
                 ):
                     if not action == other_action:
                         return False
diff --git a/stormvogel/result.py b/stormvogel/result.py
index c99e728..f8260e6 100644
--- a/stormvogel/result.py
+++ b/stormvogel/result.py
@@ -75,7 +75,7 @@ def add_scheduler(self, stormpy_scheduler: stormpy.storage.Scheduler):
             self.stormpy_scheduler = stormpy_scheduler
             taken_actions = {}
             for state in self.model.states.values():
-                taken_actions[state.id] = self.model.get_action(
+                taken_actions[state.id] = stormvogel.model.Action.create(
                     str(stormpy_scheduler.get_choice(state.id))
                 )
             self.scheduler = Scheduler(self.model, taken_actions)
diff --git a/stormvogel/simulator.py b/stormvogel/simulator.py
index edcf3bd..de4f3b6 100644
--- a/stormvogel/simulator.py
+++ b/stormvogel/simulator.py
@@ -7,7 +7,7 @@
 import stormpy.examples
 from typing import Callable
 import random
-
+from stormvogel.model import EmptyAction
 
 class Path:
     """
@@ -242,9 +242,10 @@ def simulate(
 
             # we already set the rewards for the initial state/stateaction
             if model.supports_actions():
-                reward_model.set_state_action_reward_at_id(
-                    partial_model.get_initial_state().id,
-                    model.rewards[index].get_state_reward(model.get_initial_state()),
+                reward_model.set_state_action_reward(
+                    partial_model.get_initial_state(),
+                    EmptyAction,
+                    model.rewards[index].get_state_reward(model.get_initial_state())
                 )
             else:
                 reward_model.set_state_reward(
diff --git a/stormvogel/visualization.py b/stormvogel/visualization.py
index 3889682..ef45a7f 100644
--- a/stormvogel/visualization.py
+++ b/stormvogel/visualization.py
@@ -160,7 +160,7 @@ def __add_transitions(self) -> None:
         # In the visualization, both actions and states are nodes, so we need to keep track of how many actions we already have.
         for state_id, transition in self.model.transitions.items():
             for action, branch in transition.transition.items():
-                if action.strict_eq(stormvogel.model.EmptyAction):
+                if action == stormvogel.model.EmptyAction:
                     # Only draw probabilities
                     for prob, target in branch.branch:
                         self.nt.add_edge(
@@ -175,7 +175,7 @@ def __add_transitions(self) -> None:
                         choice = self.scheduler.get_choice_of_state(
                             state=self.model.get_state_by_id(state_id)
                         )
-                        if action.strict_eq(choice):
+                        if action == choice:
                             group = "scheduled_actions"
 
                     reward = self.__format_rewards(
@@ -185,7 +185,7 @@ def __add_transitions(self) -> None:
                     # Add the action's node
                     self.nt.add_node(
                         id=action_id,
-                        label=action.name + reward,
+                        label=",".join(action.labels) + reward,
                         group=group,
                         position_dict=self.positions,
                     )
diff --git a/tests/saved_test_layout.json b/tests/saved_test_layout.json
index ef74d4a..c3802e7 100644
--- a/tests/saved_test_layout.json
+++ b/tests/saved_test_layout.json
@@ -118,4 +118,4 @@
   "init": {
     "color": "TEST_COLOR"
   }
-}
+}
\ No newline at end of file
diff --git a/tests/test_mapping.py b/tests/test_mapping.py
index 9ccc604..acf0fe9 100644
--- a/tests/test_mapping.py
+++ b/tests/test_mapping.py
@@ -139,43 +139,41 @@ def test_stormpy_to_stormvogel_and_back_mdp():
 
     assert sparse_equal(stormpy_mdp, new_stormpy_mdp)
 
+def test_stormvogel_to_stormpy_and_back_mdp():
+    # we test it for monty hall mdp
+    stormvogel_mdp = examples.monty_hall.create_monty_hall_mdp()
 
-# TODO re-introduce this test once names are removed from actions
-# def test_stormvogel_to_stormpy_and_back_mdp():
-#     # we test it for monty hall mdp
-#     stormvogel_mdp = examples.monty_hall.create_monty_hall_mdp()
-
-#     # we additionally test if reward models work
-#     rewardmodel = stormvogel_mdp.add_rewards("rewardmodel")
-#     rewardmodel.set_from_rewards_vector(list(range(67)))
-#     rewardmodel2 = stormvogel_mdp.add_rewards("rewardmodel2")
-#     rewardmodel2.set_from_rewards_vector(list(range(67)))
-
-#     # print(stormvogel_mdp)
-#     stormpy_mdp = stormvogel.mapping.stormvogel_to_stormpy(stormvogel_mdp)
-#     # print(stormpy_mdp)
-#     new_stormvogel_mdp = stormvogel.mapping.stormpy_to_stormvogel(stormpy_mdp)
-
-#     # rew0 = sorted(stormvogel_mdp.rewards)[0].rewards
-#     # rew1 = sorted(new_stormvogel_mdp.rewards)[0].rewards
-#     # for k,v in rew0.items():
-#     #     if v != rew1[k]:
-#     #         print("original:", v, "\tsv:", rew1[k])
-#     # print(rew0 == rew1)
-#     # quit()
-
-#     # print(new_stormvogel_mdp)
-#     # print(sorted(stormvogel_mdp.rewards)[0] == sorted(new_stormvogel_mdp.rewards)[0])
-#     # quit()
-#     # print(sorted(stormvogel_mdp.rewards)[0].rewards)
-#     # print()
-#     # print(sorted(new_stormvogel_mdp.rewards)[0].rewards)
-#     # print(sorted(new_stormvogel_mdp.rewards)[0].rewards == sorted(stormvogel_mdp.rewards)[0].rewards)
-#     # print(sorted(new_stormvogel_mdp.rewards)[0] == sorted(stormvogel_mdp.rewards)[0])
-#     # print(new_stormvogel_mdp == stormvogel_mdp)
-#     # quit()
-
-#     assert new_stormvogel_mdp == stormvogel_mdp
+    # we additionally test if reward models work
+    rewardmodel = stormvogel_mdp.add_rewards("rewardmodel")
+    rewardmodel.set_from_rewards_vector(list(range(67)))
+    rewardmodel2 = stormvogel_mdp.add_rewards("rewardmodel2")
+    rewardmodel2.set_from_rewards_vector(list(range(67)))
+
+    # print(stormvogel_mdp)
+    stormpy_mdp = stormvogel.mapping.stormvogel_to_stormpy(stormvogel_mdp)
+    # print(stormpy_mdp)
+    new_stormvogel_mdp = stormvogel.mapping.stormpy_to_stormvogel(stormpy_mdp)
+
+    # rew0 = sorted(stormvogel_mdp.rewards)[0].rewards
+    # rew1 = sorted(new_stormvogel_mdp.rewards)[0].rewards
+    # for k,v in rew0.items():
+    #     if v != rew1[k]:
+    #         print("original:", v, "\tsv:", rew1[k])
+    # print(rew0 == rew1)
+    # quit()
+
+    # print(new_stormvogel_mdp)
+    # print(sorted(stormvogel_mdp.rewards)[0] == sorted(new_stormvogel_mdp.rewards)[0])
+    # quit()
+    # print(sorted(stormvogel_mdp.rewards)[0].rewards)
+    # print()
+    # print(sorted(new_stormvogel_mdp.rewards)[0].rewards)
+    # print(sorted(new_stormvogel_mdp.rewards)[0].rewards == sorted(stormvogel_mdp.rewards)[0].rewards)
+    # print(sorted(new_stormvogel_mdp.rewards)[0] == sorted(stormvogel_mdp.rewards)[0])
+    # print(new_stormvogel_mdp == stormvogel_mdp)
+    # quit()
+
+    assert new_stormvogel_mdp == stormvogel_mdp
 
 
 def test_stormvogel_to_stormpy_and_back_ctmc():
@@ -211,6 +209,10 @@ def test_stormvogel_to_stormpy_and_back_pomdp():
     # print(stormpy_pomdp)
     new_stormvogel_pomdp = stormvogel.mapping.stormpy_to_stormvogel(stormpy_pomdp)
     # print(new_stormvogel_pomdp)
+    # print(stormvogel_pomdp.actions)
+    # print()
+    # print(new_stormvogel_pomdp.actions)
+    # print(stormvogel_pomdp.actions == new_stormvogel_pomdp.actions)
 
     assert new_stormvogel_pomdp == stormvogel_pomdp
 
@@ -228,7 +230,7 @@ def test_stormpy_to_stormvogel_and_back_pomdp():
     assert sparse_equal(stormpy_pomdp, new_stormpy_pomdp)
 
 
-"""
+
 def test_stormvogel_to_stormpy_and_back_ma():
     # we create a stormpy representation of an example ma
     stormvogel_ma = examples.simple_ma.create_simple_ma()
@@ -252,4 +254,4 @@ def test_stormpy_to_stormvogel_and_back_ma():
     # print(new_stormpy_ma)
 
     assert sparse_equal(stormpy_ma, new_stormpy_ma)
-"""
+
diff --git a/tests/test_model_methods.py b/tests/test_model_methods.py
index 4ab92e3..4a75e07 100644
--- a/tests/test_model_methods.py
+++ b/tests/test_model_methods.py
@@ -4,15 +4,16 @@
 import examples.nuclear_fusion_ctmc
 import pytest
 from typing import cast
+from stormvogel.model import EmptyAction
 
 
 def test_available_actions():
     mdp = examples.monty_hall.create_monty_hall_mdp()
 
     action = [
-        stormvogel.model.Action(name="open0", labels=frozenset()),
-        stormvogel.model.Action(name="open1", labels=frozenset()),
-        stormvogel.model.Action(name="open2", labels=frozenset()),
+        stormvogel.model.Action(frozenset({"open0"})),
+        stormvogel.model.Action(frozenset({"open1"})),
+        stormvogel.model.Action(frozenset({"open2"})),
     ]
     assert mdp.get_state_by_id(1).available_actions() == action
 
@@ -25,7 +26,7 @@ def test_get_outgoing_transitions():
     mdp = examples.monty_hall.create_monty_hall_mdp()
 
     transitions = mdp.get_initial_state().get_outgoing_transitions(
-        stormvogel.model.Action(name="empty", labels=frozenset())
+        stormvogel.model.Action(labels=frozenset())
     )
 
     probabilities, states = zip(*transitions)
@@ -81,7 +82,7 @@ def test_transition_from_shorthand():
     # Then we test it for a model with actions
     mdp = stormvogel.model.new_mdp()
     state = mdp.new_state()
-    action = mdp.new_action("0", frozenset({"action"}))
+    action = mdp.new_action(frozenset({"action"}))
     transition_shorthand = [(action, state)]
     branch = stormvogel.model.Branch(
         cast(list[tuple[stormvogel.model.Number, stormvogel.model.State]], [(1, state)])
@@ -347,38 +348,38 @@ def test_get_state_action_reward():
 
     assert rewardmodel.get_state_action_reward(state, action) == 5
 
+def test_set_state_reward():
+    # we create an mdp:
+    mdp = stormvogel.model.new_mdp()
+    action = stormvogel.model.Action.create()
+    mdp.add_transitions(mdp.get_initial_state(), [(action, mdp.get_initial_state())])
 
-# TODO re-introduce this test once names are removed from actions.
-# def test_set_state_action_reward():
-#     # we create an mdp:
-#     mdp = stormvogel.model.new_mdp()
-#     action = stormvogel.model.Action("0", frozenset())
-#     mdp.add_transitions(mdp.get_initial_state(), [(action, mdp.get_initial_state())])
+    # we make a reward model using the set_state_action_reward method:
+    rewardmodel = mdp.add_rewards("rewardmodel")
+    rewardmodel.set_state_action_reward(mdp.get_initial_state(), action, 5)
 
-#     # we make a reward model using the set_state_action_reward method:
-#     rewardmodel = mdp.add_rewards("rewardmodel")
-#     rewardmodel.set_state_action_reward(mdp.get_initial_state(), action, 5)
+    # we make a reward model manually:
+    other_rewardmodel = stormvogel.model.RewardModel("rewardmodel", mdp, {(0, EmptyAction): 5})
 
-#     # we make a reward model manually:
-#     other_rewardmodel = stormvogel.model.RewardModel("rewardmodel", mdp, {(0, EmptyAction): 5})
+    print(rewardmodel.rewards)
+    print()
+    print(other_rewardmodel.rewards)
 
-#     print(rewardmodel.rewards)
-#     print()
-#     print(other_rewardmodel.rewards)
-#     quit()
+    assert rewardmodel == other_rewardmodel
 
-#     assert rewardmodel == other_rewardmodel
+def test_set_state_action_reward():
 
-#     # we create an mdp:
-#     mdp = examples.monty_hall.create_monty_hall_mdp()
+    # we create an mdp:
+    mdp = examples.monty_hall.create_monty_hall_mdp()
 
-#     # we add a reward model with only one reward
-#     rewardmodel = mdp.add_rewards("rewardmodel")
-#     state = mdp.get_state_by_id(2)
-#     action = state.available_actions()[1]
-#     rewardmodel.set_state_action_reward(state, action, 3)
+    # we add a reward model with only one reward
+    rewardmodel = mdp.add_rewards("rewardmodel")
+    state = mdp.get_state_by_id(2)
+    action = state.available_actions()[0]
+    print(action)
+    rewardmodel.set_state_action_reward(state, action, 3)
 
-#     # we make a reward model manually:
-#     other_rewardmodel = stormvogel.model.RewardModel("rewardmodel", mdp, {(5, EmptyAction): 3})
+    # we make a reward model manually:
+    other_rewardmodel = stormvogel.model.RewardModel("rewardmodel", mdp, {(2, action): 3})
 
-#     assert rewardmodel == other_rewardmodel
+    assert rewardmodel == other_rewardmodel
diff --git a/tests/test_result.py b/tests/test_result.py
index 8ffecbd..4a470b1 100644
--- a/tests/test_result.py
+++ b/tests/test_result.py
@@ -635,7 +635,7 @@ def test_convert_model_checker_results_mdp():
         0,
         0,
     ] == [
-        int(list(action.name)[0])
+        int(list(action.labels)[0])
         for action in stormvogel_result.scheduler.taken_actions.values()
     ]
 
diff --git a/tests/test_simulator.py b/tests/test_simulator.py
index 95ede95..8eaaa83 100644
--- a/tests/test_simulator.py
+++ b/tests/test_simulator.py
@@ -1,99 +1,105 @@
-# TODO re-introduce tests once action names are removed.
-# def test_simulate():
-#     # we make a die dtmc and run the simulator with it
-#     dtmc = examples.die.create_die_dtmc()
-#     rewardmodel = dtmc.add_rewards("rewardmodel")
-#     for stateid in dtmc.states.keys():
-#         rewardmodel.rewards[(stateid, EmptyAction)] = 3
-#     rewardmodel2 = dtmc.add_rewards("rewardmodel2")
-#     for stateid in dtmc.states.keys():
-#         rewardmodel2.rewards[(stateid, EmptyAction)] = 2
-#     rewardmodel3 = dtmc.add_rewards("rewardmodel3")
-#     for stateid in dtmc.states.keys():
-#         rewardmodel3.rewards[(stateid, EmptyAction)] = 1
-#     partial_model = stormvogel.simulator.simulate(dtmc, runs=5, steps=1, seed=1)
-
-#     # we make the partial model that should be created by the simulator
-#     other_dtmc = stormvogel.model.new_dtmc()
-#     other_dtmc.get_initial_state().set_transitions(
-#         [
-#             (1 / 6, other_dtmc.new_state("rolled5")),
-#             (1 / 6, other_dtmc.new_state("rolled0")),
-#             (1 / 6, other_dtmc.new_state("rolled1")),
-#         ]
-#     )
-
-#     rewardmodel = other_dtmc.add_rewards("rewardmodel")
-#     for stateid in other_dtmc.states.keys():
-#         rewardmodel.rewards[(stateid, EmptyAction)] = float(3)
-#     rewardmodel2 = other_dtmc.add_rewards("rewardmodel2")
-#     for stateid in other_dtmc.states.keys():
-#         rewardmodel2.rewards[(stateid, EmptyAction)] = float(2)
-#     rewardmodel3 = other_dtmc.add_rewards("rewardmodel3")
-#     for stateid in other_dtmc.states.keys():
-#         rewardmodel3.rewards[(stateid, EmptyAction)] = float(1)
-
-#     assert partial_model == other_dtmc
-#     ######################################################################################################################
-#     # we make a monty hall mdp and run the simulator with it
-#     mdp = examples.monty_hall.create_monty_hall_mdp()
-#     rewardmodel = mdp.add_rewards("rewardmodel")
-#     for i in range(67):
-#         rewardmodel.rewards[i] = i
-#     rewardmodel2 = mdp.add_rewards("rewardmodel2")
-#     for i in range(67):
-#         rewardmodel2.rewards[i] = i
-
-#     taken_actions = {}
-#     for id, state in mdp.states.items():
-#         taken_actions[id] = state.available_actions()[0]
-#     scheduler = stormvogel.result.Scheduler(mdp, taken_actions)
-
-#     partial_model = stormvogel.simulator.simulate(
-#         mdp, runs=1, steps=3, seed=1, scheduler=scheduler
-#     )
-
-#     # we make the partial model that should be created by the simulator
-#     other_mdp = stormvogel.model.new_mdp()
-#     other_mdp.get_initial_state().set_transitions(
-#         [(1 / 3, other_mdp.new_state("carchosen"))]
-#     )
-#     other_mdp.get_state_by_id(1).set_transitions([(1, other_mdp.new_state("open"))])
-#     other_mdp.get_state_by_id(2).set_transitions(
-#         [(1, other_mdp.new_state("goatrevealed"))]
-#     )
-
-#     rewardmodel = other_mdp.add_rewards("rewardmodel")
-#     rewardmodel.rewards = {0: 0, 7: 7, 16: 16}
-#     rewardmodel2 = other_mdp.add_rewards("rewardmodel2")
-#     rewardmodel2.rewards = {0: 0, 7: 7, 16: 16}
-
-#     assert partial_model == other_mdp
-#     ######################################################################################################################
+import examples.die
+import examples.monty_hall
+import examples.nuclear_fusion_ctmc
+import examples.monty_hall_pomdp
+from stormvogel.model import EmptyAction
+import stormvogel.model
+import stormvogel.simulator
 
-#     # we test the simulator for an mdp with a lambda as Scheduler
-
-#     def scheduler(state: stormvogel.model.State) -> stormvogel.model.Action:
-#         actions = state.available_actions()
-#         return actions[0]
-
-#     mdp = examples.monty_hall.create_monty_hall_mdp()
-
-#     partial_model = stormvogel.simulator.simulate(
-#         mdp, runs=1, steps=3, seed=1, scheduler=scheduler
-#     )
-
-#     # we make the partial model that should be created by the simulator
-#     other_mdp = stormvogel.model.new_mdp()
-#     other_mdp.get_initial_state().set_transitions(
-#         [(1 / 3, other_mdp.new_state("carchosen"))]
-#     )
-#     other_mdp.get_state_by_id(1).set_transitions([(1, other_mdp.new_state("open"))])
-#     other_mdp.get_state_by_id(2).set_transitions(
-#         [(1, other_mdp.new_state("goatrevealed"))]
-#     )
-
-#     assert partial_model == other_mdp
+# TODO re-introduce tests once action names are removed.
+def test_simulate():
+    # we make a die dtmc and run the simulator with it
+    dtmc = examples.die.create_die_dtmc()
+    rewardmodel = dtmc.add_rewards("rewardmodel")
+    for stateid in dtmc.states.keys():
+        rewardmodel.rewards[(stateid, EmptyAction)] = 3
+    rewardmodel2 = dtmc.add_rewards("rewardmodel2")
+    for stateid in dtmc.states.keys():
+        rewardmodel2.rewards[(stateid, EmptyAction)] = 2
+    rewardmodel3 = dtmc.add_rewards("rewardmodel3")
+    for stateid in dtmc.states.keys():
+        rewardmodel3.rewards[(stateid, EmptyAction)] = 1
+    partial_model = stormvogel.simulator.simulate(dtmc, runs=5, steps=1, seed=1)
+
+    # we make the partial model that should be created by the simulator
+    other_dtmc = stormvogel.model.new_dtmc()
+    other_dtmc.get_initial_state().set_transitions(
+        [
+            (1 / 6, other_dtmc.new_state("rolled5")),
+            (1 / 6, other_dtmc.new_state("rolled0")),
+            (1 / 6, other_dtmc.new_state("rolled1")),
+        ]
+    )
+
+    rewardmodel = other_dtmc.add_rewards("rewardmodel")
+    for stateid in other_dtmc.states.keys():
+        rewardmodel.rewards[(stateid, EmptyAction)] = float(3)
+    rewardmodel2 = other_dtmc.add_rewards("rewardmodel2")
+    for stateid in other_dtmc.states.keys():
+        rewardmodel2.rewards[(stateid, EmptyAction)] = float(2)
+    rewardmodel3 = other_dtmc.add_rewards("rewardmodel3")
+    for stateid in other_dtmc.states.keys():
+        rewardmodel3.rewards[(stateid, EmptyAction)] = float(1)
+
+    assert partial_model == other_dtmc
+    ######################################################################################################################
+    # we make a monty hall mdp and run the simulator with it
+    mdp = examples.monty_hall.create_monty_hall_mdp()
+    rewardmodel = mdp.add_rewards("rewardmodel")
+    rewardmodel.set_from_rewards_vector(list(range(67)))
+    rewardmodel2 = mdp.add_rewards("rewardmodel2")
+    rewardmodel2.set_from_rewards_vector(list(range(67)))
+
+    taken_actions = {}
+    for id, state in mdp.states.items():
+        taken_actions[id] = state.available_actions()[0]
+    scheduler = stormvogel.result.Scheduler(mdp, taken_actions)
+
+    partial_model = stormvogel.simulator.simulate(
+        mdp, runs=1, steps=3, seed=1, scheduler=scheduler
+    )
+
+    # we make the partial model that should be created by the simulator
+    other_mdp = stormvogel.model.new_mdp()
+    other_mdp.get_initial_state().set_transitions(
+        [(1 / 3, other_mdp.new_state("carchosen"))]
+    )
+    other_mdp.get_state_by_id(1).set_transitions([(1, other_mdp.new_state("open"))])
+    other_mdp.get_state_by_id(2).set_transitions(
+        [(1, other_mdp.new_state("goatrevealed"))]
+    )
+
+    rewardmodel = other_mdp.add_rewards("rewardmodel")
+    rewardmodel.rewards = {0: 0, 7: 7, 16: 16}
+    rewardmodel2 = other_mdp.add_rewards("rewardmodel2")
+    rewardmodel2.rewards = {0: 0, 7: 7, 16: 16}
+
+    assert partial_model == other_mdp
+    ######################################################################################################################
+
+    # we test the simulator for an mdp with a lambda as Scheduler
+
+    def scheduler(state: stormvogel.model.State) -> stormvogel.model.Action:
+        actions = state.available_actions()
+        return actions[0]
+
+    mdp = examples.monty_hall.create_monty_hall_mdp()
+
+    partial_model = stormvogel.simulator.simulate(
+        mdp, runs=1, steps=3, seed=1, scheduler=scheduler
+    )
+
+    # we make the partial model that should be created by the simulator
+    other_mdp = stormvogel.model.new_mdp()
+    other_mdp.get_initial_state().set_transitions(
+        [(1 / 3, other_mdp.new_state("carchosen"))]
+    )
+    other_mdp.get_state_by_id(1).set_transitions([(1, other_mdp.new_state("open"))])
+    other_mdp.get_state_by_id(2).set_transitions(
+        [(1, other_mdp.new_state("goatrevealed"))]
+    )
+
+    assert partial_model == other_mdp
 
 
 # def test_simulate_path():
diff --git a/tests/test_visualization.py b/tests/test_visualization.py
index 6e7e36a..0001478 100644
--- a/tests/test_visualization.py
+++ b/tests/test_visualization.py
@@ -50,27 +50,25 @@ def test_show(mocker):
     MockNetwork.add_edge.assert_any_call(0, 1, label="1")
     assert MockNetwork.add_edge.call_count == 1
 
-
-# TODO re-introduce test once action names are removed.
-# def test_rewards(mocker):
-#     MockNetwork = boilerplate(mocker)
-#     model, one, init = simple_model()
-#     model.set_transitions(init, [(1, one)])
-#     model.add_rewards("LOL")
-#     model.get_rewards("LOL").set_state_reward(one, 37)
-#     model.add_rewards("HIHI")
-#     model.get_rewards("HIHI").set_state_reward(one, 42)
-#     vis = Visualization(model=model)
-#     vis.show()
-#     MockNetwork.add_node.assert_any_call(
-#         0, label="init", group="states", position_dict={}
-#     )  # type: ignore
-#     MockNetwork.add_node.assert_any_call(
-#         1, label="one\n€\tLOL: 37\tHIHI: 42", group="states", position_dict={}
-#     )  # type: ignore
-#     assert MockNetwork.add_node.call_count == 2
-#     MockNetwork.add_edge.assert_any_call(0, 1, label="1")
-#     assert MockNetwork.add_edge.call_count == 1
+def test_rewards(mocker):
+    MockNetwork = boilerplate(mocker)
+    model, one, init = simple_model()
+    model.set_transitions(init, [(1, one)])
+    model.add_rewards("LOL")
+    model.get_rewards("LOL").set_state_reward(one, 37)
+    model.add_rewards("HIHI")
+    model.get_rewards("HIHI").set_state_reward(one, 42)
+    vis = Visualization(model=model)
+    vis.show()
+    MockNetwork.add_node.assert_any_call(
+        0, label="init", group="states", position_dict={}
+    )  # type: ignore
+    MockNetwork.add_node.assert_any_call(
+        1, label="one\n€\tLOL: 37\tHIHI: 42", group="states", position_dict={}
+    )  # type: ignore
+    assert MockNetwork.add_node.call_count == 2
+    MockNetwork.add_edge.assert_any_call(0, 1, label="1")
+    assert MockNetwork.add_edge.call_count == 1
 
 
 def test_results_count(mocker):
@@ -98,8 +96,8 @@ def test_results_scheduler(mocker):
     MockNetwork = boilerplate(mocker)
     model = Model("mdp", model_type=ModelType.MDP)
     init = model.get_initial_state()
-    good = model.new_action("good", frozenset(["GOOD"]))
-    bad = model.new_action("bad", frozenset(["BAD"]))
+    good = model.new_action(frozenset(["GOOD"]))
+    bad = model.new_action(frozenset(["BAD"]))
     end = model.new_state("end")
     model.set_transitions(init, [(good, end), (bad, end)])
     scheduler = Scheduler(model, {0: good})
@@ -107,10 +105,10 @@ def test_results_scheduler(mocker):
     vis = Visualization(model=model, result=result)
     vis.show()
     MockNetwork.add_node.assert_any_call(
-        id=10000000001, label="bad", group="actions", position_dict={}
+        id=10000000001, label="BAD", group="actions", position_dict={}
     )
     MockNetwork.add_node.assert_any_call(
-        id=10000000000, label="good", group="scheduled_actions", position_dict={}
+        id=10000000000, label="GOOD", group="scheduled_actions", position_dict={}
     )
 
     assert MockNetwork.add_node.call_count == 4