AI state machine¶

a classical patrol / chase / attack / return enemy

â–¶ Run in browser

Tags: ai state-machine 2d

A finite-state machine drives an enemy through four plain-Python states: PATROL walks a waypoint loop, CHASE closes in while the player is inside the vision radius with clear line of sight (a smoke cloud blocks sight, not movement), ATTACK strikes on a cooldown at close range, and RETURN walks back to the route once the player escapes. The live state hangs above the enemy and every transition lands in an on-screen log. WASD or arrows move the player.

What it demonstrates¶

  • A hand-rolled FSM: one small class per state with enter/update hooks, and a dict of named states on the enemy. No behaviour trees, no LLM, no engine magic; this is the classical counterpoint to the LLM examples beside it.

  • Transitions as data: each update returns the next state’s name (or None), and the enemy logs every edge through a callback.

  • Perception with hysteresis: a vision radius to acquire, a larger radius to lose, and a segment-vs-rect line-of-sight test through a sight blocker.

Controls: WASD / arrow keys - move the player ESC - quit

Run: uv run python examples/features/ai/state_machine.py Headless self-check: uv run python examples/features/ai/state_machine.py –test

Source¶

  1"""AI state machine: a classical patrol / chase / attack / return enemy
  2
  3A finite-state machine drives an enemy through four plain-Python states:
  4PATROL walks a waypoint loop, CHASE closes in while the player is inside the
  5vision radius with clear line of sight (a smoke cloud blocks sight, not
  6movement), ATTACK strikes on a cooldown at close range, and RETURN walks back
  7to the route once the player escapes. The live state hangs above the enemy and
  8every transition lands in an on-screen log. WASD or arrows move the player.
  9
 10# /// simvx
 11# tags = ["ai", "state-machine", "2d"]
 12# web = { root = "StateMachineDemo", width = 960, height = 540 }
 13# ///
 14
 15## What it demonstrates
 16- A hand-rolled FSM: one small class per state with enter/update hooks, and a
 17  dict of named states on the enemy. No behaviour trees, no LLM, no engine
 18  magic; this is the classical counterpoint to the LLM examples beside it.
 19- Transitions as data: each update returns the next state's name (or None),
 20  and the enemy logs every edge through a callback.
 21- Perception with hysteresis: a vision radius to acquire, a larger radius to
 22  lose, and a segment-vs-rect line-of-sight test through a sight blocker.
 23
 24Controls:
 25  WASD / arrow keys - move the player
 26  ESC - quit
 27
 28Run: uv run python examples/features/ai/state_machine.py
 29Headless self-check: uv run python examples/features/ai/state_machine.py --test
 30"""
 31
 32from simvx.core import Input, Key, Node2D, Vec2
 33from simvx.graphics import App
 34
 35WIDTH, HEIGHT = 960, 540
 36
 37WAYPOINTS = [Vec2(140, 140), Vec2(420, 140), Vec2(420, 420), Vec2(140, 420)]
 38SMOKE = (560.0, 180.0, 120.0, 200.0)  # x, y, w, h: blocks sight, not movement
 39
 40VISION_RADIUS = 190.0  # acquire the player inside this...
 41LOSE_RADIUS = 250.0  # ...and give up beyond this (hysteresis)
 42ATTACK_RANGE = 44.0
 43ATTACK_EXIT = ATTACK_RANGE + 12.0  # hysteresis so ATTACK does not flicker
 44ATTACK_COOLDOWN = 0.8
 45FLASH_TIME = 0.18
 46ARRIVE_DIST = 4.0
 47
 48PATROL_SPEED = 90.0
 49CHASE_SPEED = 165.0
 50RETURN_SPEED = 115.0
 51PLAYER_SPEED = 220.0
 52PLAYER_R = 12.0
 53ENEMY_R = 14.0
 54
 55STATE_COLOURS = {
 56    "PATROL": (0.35, 0.8, 0.4, 1.0),
 57    "CHASE": (1.0, 0.6, 0.15, 1.0),
 58    "ATTACK": (1.0, 0.25, 0.25, 1.0),
 59    "RETURN": (0.4, 0.6, 1.0, 1.0),
 60}
 61
 62
 63def segment_crosses_rect(a: Vec2, b: Vec2, rect) -> bool:
 64    """True if the segment a-b passes through the axis-aligned rect (x, y, w, h)."""
 65    x, y, w, h = rect
 66    dx, dy = b.x - a.x, b.y - a.y
 67    t0, t1 = 0.0, 1.0
 68    for p, q in ((-dx, a.x - x), (dx, x + w - a.x), (-dy, a.y - y), (dy, y + h - a.y)):
 69        if p == 0.0:
 70            if q < 0.0:
 71                return False  # parallel and entirely outside this edge
 72        else:
 73            t = q / p
 74            if p < 0.0:
 75                t0 = max(t0, t)
 76            else:
 77                t1 = min(t1, t)
 78            if t0 > t1:
 79                return False
 80    return True
 81
 82
 83class State:
 84    """One FSM state. update() returns the next state's name, or None to stay."""
 85
 86    name = "?"
 87
 88    def enter(self, e: Enemy) -> None:
 89        pass
 90
 91    def update(self, e: Enemy, dt: float) -> str | None:
 92        return None
 93
 94
 95class Patrol(State):
 96    name = "PATROL"
 97
 98    def update(self, e, dt):
 99        if e.sees_player():
100            return "CHASE"
101        target = e.waypoints[e.wp_index]
102        e.pos = e.pos.move_toward(target, PATROL_SPEED * dt)
103        if e.pos.distance_to(target) < ARRIVE_DIST:
104            e.wp_index = (e.wp_index + 1) % len(e.waypoints)
105        return None
106
107
108class Chase(State):
109    name = "CHASE"
110
111    def enter(self, e):
112        self.lost = 0.0  # seconds since line of sight was last clear
113
114    def update(self, e, dt):
115        dist = e.pos.distance_to(e.player_pos)
116        if dist <= ATTACK_RANGE:
117            return "ATTACK"
118        self.lost = 0.0 if e.has_line_of_sight() else self.lost + dt
119        if dist > LOSE_RADIUS or self.lost > 1.5:
120            return "RETURN"
121        e.pos = e.pos.move_toward(e.player_pos, CHASE_SPEED * dt)
122        return None
123
124
125class Attack(State):
126    name = "ATTACK"
127
128    def enter(self, e):
129        e.attack_timer = 0.0  # first strike lands immediately
130
131    def update(self, e, dt):
132        if e.pos.distance_to(e.player_pos) > ATTACK_EXIT:
133            return "CHASE"
134        e.attack_timer -= dt
135        if e.attack_timer <= 0.0:
136            e.hits += 1
137            e.flash = FLASH_TIME
138            e.attack_timer = ATTACK_COOLDOWN
139        return None
140
141
142class Return(State):
143    name = "RETURN"
144
145    def enter(self, e):
146        # Head for the nearest waypoint and rejoin the loop there.
147        e.wp_index = min(range(len(e.waypoints)), key=lambda i: e.pos.distance_to(e.waypoints[i]))
148
149    def update(self, e, dt):
150        if e.sees_player():
151            return "CHASE"
152        target = e.waypoints[e.wp_index]
153        e.pos = e.pos.move_toward(target, RETURN_SPEED * dt)
154        if e.pos.distance_to(target) < ARRIVE_DIST:
155            return "PATROL"
156        return None
157
158
159class Enemy:
160    """Pure-Python FSM agent: no node, no window, fully testable headless."""
161
162    def __init__(self, waypoints, on_transition=None):
163        self.waypoints = [Vec2(w.x, w.y) for w in waypoints]
164        self.pos = Vec2(self.waypoints[0].x, self.waypoints[0].y)
165        self.wp_index = 1
166        self.player_pos = Vec2(0, 0)
167        self.hits = 0
168        self.attack_timer = 0.0
169        self.flash = 0.0
170        self._on_transition = on_transition
171        self.states: dict[str, State] = {s.name: s for s in (Patrol(), Chase(), Attack(), Return())}
172        self.state = "PATROL"
173        self.states[self.state].enter(self)
174
175    def has_line_of_sight(self) -> bool:
176        return not segment_crosses_rect(self.pos, self.player_pos, SMOKE)
177
178    def sees_player(self) -> bool:
179        return self.pos.distance_to(self.player_pos) <= VISION_RADIUS and self.has_line_of_sight()
180
181    def set_state(self, name: str) -> None:
182        if name == self.state:
183            return
184        if self._on_transition is not None:
185            self._on_transition(self.state, name)
186        self.state = name
187        self.states[name].enter(self)
188
189    def update(self, dt: float, player_pos: Vec2) -> None:
190        self.player_pos = player_pos
191        self.flash = max(0.0, self.flash - dt)
192        nxt = self.states[self.state].update(self, dt)
193        if nxt is not None:
194            self.set_state(nxt)
195
196
197class StateMachineDemo(Node2D):
198    """WASD player dot versus a four-state patrol enemy."""
199
200    dynamic = True  # everything moves every frame
201
202    input_actions = {
203        "move_left": [Key.A, Key.LEFT],
204        "move_right": [Key.D, Key.RIGHT],
205        "move_up": [Key.W, Key.UP],
206        "move_down": [Key.S, Key.DOWN],
207        "quit": [Key.ESCAPE],
208    }
209
210    def on_ready(self):
211        self._player = Vec2(790, 420)
212        self._time = 0.0
213        self._log: list[str] = []
214        self._enemy = Enemy(WAYPOINTS, on_transition=self._log_edge)
215
216    def _log_edge(self, frm: str, to: str) -> None:
217        self._log.append(f"{self._time:5.1f}s  {frm} -> {to}")
218        del self._log[:-8]
219
220    def on_update(self, dt: float):
221        if Input.is_action_just_pressed("quit"):
222            self.app.quit()
223        self._time += dt
224
225        move = Input.get_vector("move_left", "move_right", "move_up", "move_down")
226        self._player = (self._player + move * (PLAYER_SPEED * dt)).clamped(
227            Vec2(PLAYER_R, PLAYER_R), Vec2(WIDTH - PLAYER_R, HEIGHT - PLAYER_R)
228        )
229        self._enemy.update(dt, self._player)
230
231    def on_draw(self, renderer):
232        e = self._enemy
233        colour = STATE_COLOURS[e.state]
234
235        # Smoke cloud: a sight blocker the FSM's line-of-sight test respects.
236        renderer.draw_rect((SMOKE[0], SMOKE[1]), (SMOKE[2], SMOKE[3]), colour=(0.55, 0.55, 0.6, 0.35), filled=True)
237        renderer.draw_rect((SMOKE[0], SMOKE[1]), (SMOKE[2], SMOKE[3]), colour=(0.7, 0.7, 0.75, 0.6))
238        renderer.draw_text("smoke", (SMOKE[0] + 8, SMOKE[1] + 8), colour=(0.75, 0.75, 0.8))
239
240        # Patrol route, with the current target waypoint highlighted.
241        renderer.draw_lines([(w.x, w.y) for w in e.waypoints], closed=True, colour=(0.4, 0.45, 0.55, 0.8))
242        for i, w in enumerate(e.waypoints):
243            active = i == e.wp_index and e.state in ("PATROL", "RETURN")
244            wp_colour = (0.9, 0.9, 0.5, 1.0) if active else (0.5, 0.55, 0.65, 1.0)
245            renderer.draw_circle((w.x, w.y), 6, colour=wp_colour, filled=True)
246
247        # Perception: vision circle in the state colour, attack range up close.
248        renderer.draw_circle((e.pos.x, e.pos.y), VISION_RADIUS, colour=(*colour[:3], 0.25), segments=64)
249        if e.state in ("CHASE", "ATTACK"):
250            renderer.draw_circle((e.pos.x, e.pos.y), ATTACK_RANGE, colour=(1.0, 0.4, 0.4, 0.5))
251            los_ok = e.has_line_of_sight()
252            renderer.draw_line(
253                (e.pos.x, e.pos.y),
254                (self._player.x, self._player.y),
255                colour=(1.0, 0.8, 0.3, 0.5) if los_ok else (0.5, 0.5, 0.5, 0.35),
256            )
257
258        # The enemy, its attack flash, and the state label above it.
259        if e.flash > 0.0:
260            ring = ENEMY_R + 14 * (1 - e.flash / FLASH_TIME)
261            renderer.draw_circle((e.pos.x, e.pos.y), ring, colour=(1.0, 0.3, 0.2, 0.9))
262        renderer.draw_circle((e.pos.x, e.pos.y), ENEMY_R, colour=colour, filled=True)
263        renderer.draw_text(e.state, (e.pos.x, e.pos.y - ENEMY_R - 24), colour=colour, alignment="centre", scale=1.4)
264
265        # The player.
266        renderer.draw_circle((self._player.x, self._player.y), PLAYER_R, colour=(0.3, 0.85, 1.0, 1.0), filled=True)
267
268        # Transition log, newest at the bottom.
269        renderer.draw_text("Transitions", (WIDTH - 250, 10), colour=(0.85, 0.85, 0.9), scale=1.2)
270        for i, line in enumerate(self._log):
271            renderer.draw_text(line, (WIDTH - 250, 34 + i * 18), colour=(0.65, 0.7, 0.75))
272
273        # HUD.
274        renderer.draw_text("AI State Machine", (10, 10), colour=(1.0, 1.0, 1.0), scale=2)
275        renderer.draw_text(f"Hits taken: {e.hits}", (10, 44), colour=(1.0, 0.6, 0.6))
276        renderer.draw_text("WASD / arrows: move   ESC: quit", (10, HEIGHT - 28), colour=(0.6, 0.6, 0.6))
277
278
279def _selftest() -> bool:
280    """Logic-level: drive the FSM directly and check every edge, no window."""
281    dt = 1.0 / 60.0
282    ok = True
283
284    def check(label: str, passed: bool, detail: str) -> None:
285        nonlocal ok
286        ok = ok and passed
287        print(f"{'ok  ' if passed else 'FAIL'} {label}: {detail}")
288
289    # Line of sight: through the smoke is blocked, beside it is clear.
290    check(
291        "a segment through the smoke is blocked",
292        segment_crosses_rect(Vec2(520, 280), Vec2(700, 280), SMOKE),
293        "horizontal ray at y=280 crosses the rect",
294    )
295    check(
296        "a segment beside the smoke is clear",
297        not segment_crosses_rect(Vec2(520, 120), Vec2(520, 460), SMOKE),
298        "vertical ray at x=520 misses the rect",
299    )
300
301    e = Enemy(WAYPOINTS)
302    e.pos, e.player_pos = Vec2(520, 280), Vec2(700, 280)
303    check(
304        "smoke hides a player inside the vision radius",
305        not e.sees_player(),
306        f"dist {e.pos.distance_to(e.player_pos):.0f} <= {VISION_RADIUS:.0f} but blocked",
307    )
308    e.player_pos = Vec2(520, 120)
309    check("the same distance with clear sight is seen", e.sees_player(), "no blocker on this ray")
310
311    # Undisturbed patrol: the enemy walks the loop and never leaves PATROL.
312    edges: list[tuple[str, str]] = []
313    e = Enemy(WAYPOINTS, on_transition=lambda a, b: edges.append((a, b)))
314    far = Vec2(920, 60)
315    visited = {e.wp_index}
316    for _ in range(int(20.0 / dt)):
317        e.update(dt, far)
318        visited.add(e.wp_index)
319    check(
320        "a distant player leaves the enemy in PATROL",
321        e.state == "PATROL" and not edges,
322        f"state {e.state}, {len(edges)} edges",
323    )
324    check("the patrol visits every waypoint", visited == set(range(len(WAYPOINTS))), f"visited {sorted(visited)}")
325
326    # The full loop of edges: acquire, close, strike, lose, walk home.
327    edges.clear()
328    near = Vec2(WAYPOINTS[0].x + 100, WAYPOINTS[0].y)
329    for _ in range(int(6.0 / dt)):
330        e.update(dt, near)
331    check("the enemy reaches ATTACK on a nearby player", e.state == "ATTACK", f"state {e.state}")
332    check("attacks land on the cooldown", e.hits >= 2, f"{e.hits} hits in 6s at {ATTACK_COOLDOWN}s cooldown")
333    hits_before = e.hits
334    for _ in range(int(20.0 / dt)):
335        e.update(dt, far)
336    check("an escaped player sends the enemy home", e.state == "PATROL", f"state {e.state}")
337    check("no attacks land once the player escapes", e.hits == hits_before, f"{e.hits - hits_before} extra hits")
338    expected = [
339        ("PATROL", "CHASE"),
340        ("CHASE", "ATTACK"),
341        ("ATTACK", "CHASE"),
342        ("CHASE", "RETURN"),
343        ("RETURN", "PATROL"),
344    ]
345    check("the log records exactly the expected edges", edges == expected, f"{edges}")
346
347    print("SELFTEST:", "PASS" if ok else "FAIL")
348    return ok
349
350
351if __name__ == "__main__":
352    import sys
353
354    if "--test" in sys.argv:
355        sys.exit(0 if _selftest() else 1)
356    App(title="AI State Machine", width=WIDTH, height=HEIGHT).run(StateMachineDemo())