Babysnake training works

babysnake/__init__.py

@@ -9,6 +9,14 @@ The game ends when energy reaches 0.
 from random import randint, choice
 from retro.game import Game
 
+ACTIONS = ["KEY_RIGHT", "KEY_DOWN", "KEY_LEFT", "KEY_UP"]
+
+
+def get_state(game):
+    """Return the game state as a hashable (agent_x, agent_y, food_x, food_y) tuple."""
+    s = game.state
+    return (int(s['agent_x']), int(s['agent_y']), int(s['food_x']), int(s['food_y']))
+
 BOARD_SIZE = 4
 START_ENERGY = 50
 FOOD_ENERGY = 30

babysnake/pyproject.toml

@@ -1,4 +1,4 @@
 [tool.retro-gamer]
 actions = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"]
 reward = "reward"
-observation_function = "babysnake_env:get_state"
+observation_function = "babysnake:get_state"