fix examples

Author: randriu

README.md

@@ -24,7 +24,7 @@ ats.set_initial_states([ats.num_states - 1])
 umbi.ats.write(ats, "out.umb")
 ```
 
-More examples can be found in the [./examples](./examples) folder.
+More examples can be found in the [examples](./examples) folder.
 
 ## API
 
@@ -34,7 +34,7 @@ More examples can be found in the [./examples](./examples) folder.
 
 **[`ExplicitUmb`](umbi/umb/explicit_umb.py)** - in-memory representation of a typical umbfile. Attributes are standard Python objects (lists, dicts, dataclasses) providing a deserialized view of the file contents.
 
-**[`ExplicitAts`](umbi/ats/explicit_ats.py)** - format-agnostic abstraction for annotated transition systems (states, transitions, annotations). Recommended for most use cases: easiest to use programmatically and remains stable across UMB format changes.
+**[`ExplicitAts`](umbi/ats/explicit_ats.py)** - format-agnostic abstraction for annotated transition systems (states, transitions, annotations). Recommended for most use cases: easiest to use programmatically and remains stable across UMB format changes. See [umbi.ats.examples](./umbi/ats/examples/) for usage examples.
 
 ## CLI
 

umbi/ats/examples/__init__.py

@@ -1,11 +1,12 @@
 """ATS model examples."""
 
-from .grid import grid_ats
-from .random_game import random_game_ats
-from .random_walk import random_walk_ats
+from .grid import grid, random_grid_string
+from .random_game import random_game
+from .random_walk import random_walk
 
 __all__ = [
-    "grid_ats",
-    "random_game_ats",
-    "random_walk_ats",
+    "random_grid_string",
+    "grid",
+    "random_game",
+    "random_walk",
 ]

umbi/ats/examples/grid.py

@@ -1,6 +1,7 @@
 """umbi demo: Create an ExplicitAts from a grid string."""
 
 import logging
+import random
 from fractions import Fraction
 
 from ..explicit_ats import ExplicitAts, TimeType
@@ -9,7 +10,48 @@
 logger = logging.getLogger(__name__)
 
 
-def grid_ats(grid: str) -> ExplicitAts:
+def random_grid_string(width: int, height: int, seed: int | None = None) -> str:
+    """
+    Generate a random grid string of size width x height.
+
+    :param width: the width of the grid
+    :param height: the height of the grid
+    :param seed: optional seed for reproducible randomness
+    :return: A grid string with:
+        - exactly 1 'i' (initial state)
+        - random number of 'g' (goal states)
+        - random number of 'x' (obstacles)
+        - rest are '.' (empty cells)
+    """
+    if seed is not None:
+        random.seed(seed)
+
+    total_cells = width * height
+    cells = ["."] * total_cells
+
+    # Place exactly one 'i'
+    i_idx = random.randint(0, total_cells - 1)
+    cells[i_idx] = "i"
+
+    # Random number of 'g's (at least 1, up to ~5% of remaining cells)
+    num_goals = random.randint(1, max(1, total_cells // 20))
+    goal_positions = random.sample([i for i in range(total_cells) if i != i_idx], min(num_goals, total_cells - 1))
+    for pos in goal_positions:
+        cells[pos] = "g"
+
+    # Random number of 'x's from remaining cells (20 to ~50% of remaining cells)
+    remaining_positions = [i for i in range(total_cells) if cells[i] == "."]
+    num_obstacles = random.randint(max(1, len(remaining_positions) // 5), max(1, len(remaining_positions) // 2))
+    obstacle_positions = random.sample(remaining_positions, min(num_obstacles, len(remaining_positions)))
+    for pos in obstacle_positions:
+        cells[pos] = "x"
+
+    # Convert to grid string (width x height, with newlines)
+    grid_rows = ["".join(cells[i * width : (i + 1) * width]) for i in range(height)]
+    return "\n".join(reversed(grid_rows))  # reverse for top-to-bottom display
+
+
+def grid(grid: str) -> ExplicitAts:
     """
     Create a simple ATS from a rectangular grid string.
 
@@ -52,6 +94,8 @@ def grid_ats(grid: str) -> ExplicitAts:
             elif c == "g":
                 goal_states.add(state)
 
+    num_states = len(cell_to_state)
+    ats.add_states(num_states)
     ats.set_initial_states(list(initial_states))
 
     if ats.variable_valuations is None:
@@ -76,7 +120,6 @@ def grid_ats(grid: str) -> ExplicitAts:
     ats.num_choice_actions = len(ats.choice_action_to_name)
 
     for (x, y), state in cell_to_state.items():
-        ats.add_state()
         for direction, (dx, dy) in direction_dxdy.items():
             target = (x + dx, y + dy)
             choice = ats.add_state_choice(state=state)
@@ -90,10 +133,10 @@ def grid_ats(grid: str) -> ExplicitAts:
                 choice.add_branch(target=state)
 
     annotation = ats.add_ap_annotation(name="goal")
-    annotation.set_state_values([s in goal_states for s in range(ats.num_states)])
+    annotation.set_state_values([s in goal_states for s in ats.states])
 
     annotation = ats.add_reward_annotation(name="step_cost")
-    annotation.set_choice_values([1 for _ in range(ats.num_choices)])
+    annotation.set_choice_values([1 for _ in ats.choices])
 
     ats.validate()
     return ats

umbi/ats/examples/random_game.py

@@ -173,7 +173,7 @@ def random_transition_function(
     return delta, actions_at_state, all_actions
 
 
-def random_game_ats(num_states: int, seed: int | None = None) -> ExplicitAts:
+def random_game(num_states: int, seed: int | None = None) -> ExplicitAts:
     """
     Generate a random stochastic game connected from initial state.
 
@@ -196,9 +196,6 @@ def random_game_ats(num_states: int, seed: int | None = None) -> ExplicitAts:
     # set player information
     ats.state_to_player = [0 if is_player0 else 1 for is_player0 in is_player0_state]
 
-    # set initial state to the state with index 0
-    ats.set_initial_states([0])
-
     # build the ATS structure from delta and actions_at_state
     ats.num_choice_actions = len(all_actions)
     ats.choice_action_to_name = all_actions
@@ -207,10 +204,13 @@ def random_game_ats(num_states: int, seed: int | None = None) -> ExplicitAts:
     logger.info("building CSR structures ...")
     build_start = time.time()
 
-    for s in range(num_states):
-        state = ats.add_state()
+    ats.add_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=state)
+            choice = ats.add_state_choice(state=s)
             choice.action = action
 
             # get all branches for this (s, action) pair

umbi/ats/examples/random_walk.py

@@ -8,39 +8,39 @@
 logger = logging.getLogger(__name__)
 
 
-def random_walk_ats(num_states: int) -> ExplicitAts:
+def random_walk(num_states: int) -> ExplicitAts:
     # create ATS
     ats = ExplicitAts()
 
     # basic parameters
     ats.time = TimeType.DISCRETE
     ats.num_players = 0
-    ats.num_initial_states = 1
-    ats.set_initial_states([num_states // 2])  # start in the middle state
 
     # 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"]
 
+    ats.add_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):
-        ats.add_state()
-
         # left action
         choice = ats.add_state_choice(state=state)
+        choice.action = 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
         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))
 
-    ats.state_is_markovian = [True] * ats.num_states
-    ats.state_to_exit_rate = [1] * ats.num_states  # type: ignore[assignment]
-
     # example: APs
     annotation = ats.add_ap_annotation(
         name="is_terminal", alias="Terminal State", description="Indicates whether the state is terminal."

umbi/ats/explicit_ats.py

@@ -230,6 +230,14 @@ def add_state(self) -> int:
         self._state_choices.append([])
         return new_state
 
+    def add_states(self, num: int) -> list[int]:
+        """Add the given number of new states and return their indices."""
+        new_states = []
+        for _ in range(num):
+            new_state = self.add_state()
+            new_states.append(new_state)
+        return new_states
+
     def remove_state(self, state: int):
         """Remove the given state."""
         self._check_state(state)
@@ -505,7 +513,7 @@ def validate(self):
                 }[entity_class]
                 if not valuations.num_entities == num_entities:
                     raise ValueError(
-                        f"{entity_class} valuations has {valuations.num_entities} entities, expected {num_entities}"
+                        f"{entity_class.value} valuations has {valuations.num_entities} {entity_class.value}, expected {num_entities}"
                     )
                 valuations.validate()
 

umbi/ats/variable_valuations.py

@@ -382,6 +382,9 @@ def observation_valuations(self) -> EntityValuations:
     def player_valuations(self) -> EntityValuations:
         return self.get_valuations_for(EntityClass.PLAYERS)
 
+    def add_state_valuation(self):
+        return self.get_valuations_for(EntityClass.STATES)
+
     def set_state_valuations(self, values: EntityValuations):
         self.set_valuations_for(EntityClass.STATES, values)