introduce EntityToAction

Author: randriu

umbi/ats/ats_to_umb.py

@@ -10,7 +10,6 @@
 from .entity_class import EntityClass
 from .explicit_ats import ExplicitAts, TimeType
 from .model_info import ModelInfo
-from .variable_valuations import EntityClassValuations, EntityValuations
 
 
 def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
@@ -36,13 +35,10 @@ def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
     ts = umb.index.transition_system
     ats.time = TimeType(ts.time)
     ats.num_players = ts.num_players
-    ats.num_initial_states = ts.num_initial_states
-    ats.num_choice_actions = ts.num_choice_actions
-    ats.num_branch_actions = ts.num_branch_actions
     ats.player_to_name = ts.player_to_name
 
     ## values
-    ats.state_is_initial = umb.state_is_initial
+    ats.set_state_is_initial(umb.state_is_initial)
     ats.state_to_player = umb.state_to_player
 
     ats.state_is_markovian = umb.state_is_markovian
@@ -55,6 +51,20 @@ def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
     state_to_choices = (
         umb.state_to_choices if umb.state_to_choices is not None else [s for s in range(ats.num_states + 1)]
     )
+    have_choice_actions = ts.num_choice_actions > 0
+    if have_choice_actions:
+        assert umb.choice_to_choice_action is not None, "num_choice_actions > 0 but choice_to_choice_action is None"
+        ats.add_choice_actions(num_actions=ts.num_choice_actions)
+        if umb.choice_action_to_string is not None:
+            ats.choice_actions.set_names(umb.choice_action_to_string)
+    have_branch_actions = ts.num_branch_actions > 0
+    if have_branch_actions:
+        assert umb.branch_to_branch_action is not None, "num_branch_actions > 0 but branch_to_branch_action is None"
+        ats.add_branch_actions(num_actions=ts.num_branch_actions)
+        if umb.branch_action_to_string is not None:
+            ats.branch_actions.set_names(umb.branch_action_to_string)
+
+    branch_index = 0
     if ts.num_choices > 0:
         for state in ats.states:
             for choice_idx in range(state_to_choices[state], state_to_choices[state + 1]):
@@ -65,18 +75,12 @@ def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
                         assert umb.branch_to_target is not None, "num_branches > 0 but branch_to_target is None"
                         target = umb.branch_to_target[branch_idx]
                         prob = umb.branch_to_probability[branch_idx] if umb.branch_to_probability is not None else None
-                        branch_action = (
-                            umb.branch_to_branch_action[branch_idx] if umb.branch_to_branch_action is not None else None
-                        )
-                        choice.add_branch(target=target, prob=prob, action=branch_action)
-                if ts.num_choice_actions > 0:
-                    assert umb.choice_to_choice_action is not None, (
-                        "num_choice_actions > 0 but choice_to_choice_action is None"
-                    )
-                    choice.action = umb.choice_to_choice_action[choice_idx]
-
-    ats.choice_action_to_name = umb.choice_action_to_string
-    ats.branch_action_to_name = umb.branch_action_to_string
+                        choice.add_branch(target=target, prob=prob)
+                        branch_index += 1  # TODO make ats.add_branch return the branch index
+                        if have_branch_actions:
+                            ats.branch_actions[branch_index] = umb.branch_to_branch_action[branch_idx]  # type: ignore[subscript]
+                if have_choice_actions:
+                    ats.choice_actions[choice.index] = umb.choice_to_choice_action[choice_idx]  # type: ignore[subscript]
 
     # load annotations
     if umb.index.annotations is not None:
@@ -102,7 +106,7 @@ def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
     # load valuations
     if umb.index.valuations is not None:
         assert umb.valuations is not None
-        ats.variable_valuations = EntityClassValuations()
+        ats.add_variable_valuations()
         for applies_to, valuation_description in umb.index.valuations.items():
             # ignore unique and num_strings
             # assume a single valuation class
@@ -112,7 +116,7 @@ def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
 
             entity_to_valuation = umb.valuations[applies_to]
             entity_class = EntityClass(applies_to)
-            entity_valuations = EntityValuations()
+            entity_valuations = ats.variable_valuations.add_valuations_for(entity_class)
             for attribute in struct_type.attributes:
                 if attribute.lower is not None or attribute.upper is not None:
                     raise NotImplementedError("bounds on valuation variables not supported yet")
@@ -122,7 +126,6 @@ def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
             for entity, valuation in enumerate(entity_to_valuation):
                 valuation = {entity_valuations.get_variable(var_name): value for var_name, value in valuation.items()}
                 entity_valuations.set_entity_valuation(entity, valuation)
-            ats.variable_valuations.set_valuations_for(entity_class, entity_valuations)
 
     # load observations
     if ts.observations_apply_to is not None:
@@ -167,9 +170,9 @@ def explicit_ats_to_explicit_umb(ats: ExplicitAts) -> umbi.umb.ExplicitUmb:
             num_states=ats.num_states,
             num_initial_states=ats.num_initial_states,
             num_choices=ats.num_choices,
-            num_choice_actions=ats.num_choice_actions,
+            num_choice_actions=0,  # for now
             num_branches=ats.num_branches,
-            num_branch_actions=ats.num_branch_actions,
+            num_branch_actions=0,  # for now
             num_observations=ats.num_observations,
             observations_apply_to=None,  # for now
             branch_probability_type=None,  # for now
@@ -185,7 +188,7 @@ def explicit_ats_to_explicit_umb(ats: ExplicitAts) -> umbi.umb.ExplicitUmb:
     ## values
     # warning: this does not copy the values, but directly uses the lists from ats.
     # this is fine as long as we don't modify them
-    umb.state_is_initial = ats.state_is_initial
+    umb.state_is_initial = list(ats.state_is_initial)
     umb.state_to_player = ats.state_to_player
 
     umb.state_is_markovian = ats.state_is_markovian
@@ -204,8 +207,11 @@ def explicit_ats_to_explicit_umb(ats: ExplicitAts) -> umbi.umb.ExplicitUmb:
     umb.state_to_choices.append(num_choices)
 
     if ats.num_choices > 0:
-        if ats.num_choice_actions > 0:
-            umb.choice_to_choice_action = [choice.action for choice in ats.choices]  # type: ignore
+        if ats.has_choice_actions > 0:
+            umb.index.transition_system.num_choice_actions = ats.choice_actions.num_actions
+            umb.choice_to_choice_action = list(ats.choice_actions.values)
+            if ats.choice_actions.has_names:
+                umb.choice_action_to_string = list(ats.choice_actions.names)
 
         branches = []
         umb.choice_to_branches = []
@@ -223,14 +229,11 @@ def explicit_ats_to_explicit_umb(ats: ExplicitAts) -> umbi.umb.ExplicitUmb:
                 # assert isinstance(target_type, umbi.datatypes.NumericType), "branch probabilities must be numeric"
                 umb.index.transition_system.branch_probability_type = SizedType.for_type(target_type)
                 umb.branch_to_probability = vector  # type: ignore
-            if ats.num_branch_actions > 0:
-                branch_actions = [branch.action for branch in branches]
-                if any(action is None for action in branch_actions):
-                    raise ValueError("if num_branch_actions > 0, all branches must have an action")
-                umb.branch_to_branch_action = branch_actions  # type: ignore
-
-    umb.choice_action_to_string = ats.choice_action_to_name
-    umb.branch_action_to_string = ats.branch_action_to_name
+            if ats.has_branch_actions:
+                umb.index.transition_system.num_branch_actions = ats.branch_actions.num_actions
+                umb.branch_to_branch_action = list(ats.branch_actions.values)
+                if ats.branch_actions.has_names:
+                    umb.branch_action_to_string = list(ats.branch_actions.names)
 
     # add annotations
     umb.index.annotations = {}

umbi/ats/examples/grid.py

@@ -5,7 +5,6 @@
 from fractions import Fraction
 
 from ..explicit_ats import ExplicitAts, TimeType
-from ..variable_valuations import EntityClassValuations, EntityValuations
 
 logger = logging.getLogger(__name__)
 
@@ -93,7 +92,7 @@ def grid(grid: str) -> ExplicitAts:
                 goal_states.add(state)
 
     num_states = len(cell_to_state)
-    ats.add_states(num_states)
+    ats.set_num_states(num_states)
     ats.set_initial_states(list(initial_states))
 
     variable_valuations = ats.add_variable_valuations()
@@ -111,14 +110,14 @@ def grid(grid: str) -> ExplicitAts:
         "right": (1, 0),
     }
 
-    ats.choice_action_to_name = list(direction_dxdy.keys())
-    ats.num_choice_actions = len(ats.choice_action_to_name)
+    choice_actions = ats.add_choice_actions(num_actions=len(direction_dxdy))
+    choice_actions.set_names(list(direction_dxdy.keys()))
 
     for (x, y), state in cell_to_state.items():
         for direction, (dx, dy) in direction_dxdy.items():
             target = (x + dx, y + dy)
             choice = ats.add_state_choice(state=state)
-            choice.action = ats.choice_action_to_name.index(direction)
+            choice_actions[choice.index] = choice_actions.names.index(direction)
 
             if target in cell_to_state:
                 target_state = cell_to_state[target]

umbi/ats/examples/random_game.py

@@ -197,21 +197,21 @@ def random_game(num_states: int, seed: int | None = None) -> ExplicitAts:
     ats.state_to_player = [0 if is_player0 else 1 for is_player0 in is_player0_state]
 
     # build the ATS structure from delta and actions_at_state
-    ats.num_choice_actions = len(all_actions)
-    ats.choice_action_to_name = all_actions
+    choice_actions = ats.add_choice_actions(num_actions=len(all_actions))
+    choice_actions.set_names(all_actions)
 
     # build CSR structures
     logger.info("building CSR structures ...")
     build_start = time.time()
 
-    ats.add_states(num_states)
+    ats.set_num_states(num_states)
     # set initial state to the state with index 0
     ats.set_initial_states([0])
 
     for s in ats.states:
         for action in sorted(actions_at_state[s]):
             choice = ats.add_state_choice(state=s)
-            choice.action = action
+            choice_actions[choice.index] = action
 
             # get all branches for this (s, action) pair
             branches = delta[s][action]

umbi/ats/examples/random_walk.py

@@ -13,30 +13,30 @@ def random_walk(num_states: int) -> ExplicitAts:
     ats = ExplicitAts()
 
     # basic parameters
-    ats.time = TimeType.DISCRETE
+    ats.time = TimeType.STOCHASTIC
     ats.num_players = 0
 
     # two actions: left (0) and right (1), each with 0.9 success prob
-    ats.num_choice_actions = 2
-    ats.choice_action_to_name = ["left", "right"]
+    choice_actions = ats.add_choice_actions(num_actions=2)
+    choice_actions.set_names(["left", "right"])
 
-    ats.add_states(num_states)
+    ats.set_num_states(num_states)
     ats.set_initial_states([num_states // 2])  # start in the middle state
     ats.state_is_markovian = [True] * ats.num_states
     ats.state_to_exit_rate = [1] * ats.num_states  # type: ignore[assignment]
 
     # build structure
-    for state in range(num_states):
+    for state in ats.states:
         # left action
         choice = ats.add_state_choice(state=state)
-        choice.action = 0
+        choice_actions[choice.index] = 0
         left = max(0, state - 1)
         choice.add_branch(target=left, prob=Fraction(9, 10))
         choice.add_branch(target=state, prob=Fraction(1, 10))
 
         # right action
         choice = ats.add_state_choice(state=state)
-        choice.action = 1
+        choice_actions[choice.index] = 1
         right = min(num_states - 1, state + 1)
         choice.add_branch(target=right, prob=Fraction(9, 10))
         choice.add_branch(target=state, prob=Fraction(1, 10))
@@ -54,12 +54,10 @@ def random_walk(num_states: int) -> ExplicitAts:
     # wall hit penalty
     choice_penalty = [0] * ats.num_choices
     for choice in ats.get_state_choices(0):
-        assert choice.action is not None
-        if ats.choice_action_to_name[choice.action] == "left":  # left action in first state
+        if choice_actions.names[choice_actions[choice.index]] == "left":  # left action in first state
             choice_penalty[choice.index] = -10
     for choice in ats.get_state_choices(ats.num_states - 1):
-        assert choice.action is not None
-        if ats.choice_action_to_name[choice.action] == "right":  # right action in last state
+        if choice_actions.names[choice_actions[choice.index]] == "right":  # right action in last state
             choice_penalty[choice.index] = -10
     annotation = ats.add_reward_annotation(
         name="wall_hit_penalty",