Initial commit

babysnake/__init__.py

@@ -0,0 +1,117 @@
+"""BabySnake: a 4×4 grid game where an agent collects food.
+
+State: (agent_x, agent_y, food_x, food_y) — four integers.
+The agent starts with 50 energy. Each step costs 1 energy.
+Collecting food restores 30 energy and adds 1 to the score.
+The game ends when energy reaches 0.
+"""
+
+from random import randint, choice
+from retro.game import Game
+
+BOARD_SIZE = 4
+START_ENERGY = 50
+FOOD_ENERGY = 30
+
+class Forager:
+    """The player agent. Controlled with arrow keys."""
+    name = "Forager"
+    character = '@'
+    color = "green_on_black"
+    position = (0, 0)
+
+    UP = (0, -1)
+    DOWN = (0, 1)
+    LEFT = (-1, 0)
+    RIGHT = (1, 0)
+
+    def __init__(self):
+        self._direction = self.RIGHT
+
+    def handle_keystroke(self, keystroke, game):
+        if keystroke.name == "KEY_RIGHT":
+            self._direction = self.RIGHT
+        elif keystroke.name == "KEY_UP":
+            self._direction = self.UP
+        elif keystroke.name == "KEY_LEFT":
+            self._direction = self.LEFT
+        elif keystroke.name == "KEY_DOWN":
+            self._direction = self.DOWN
+
+    def play_turn(self, game):
+        x, y = self.position
+        dx, dy = self._direction
+        new_pos = (x + dx, y + dy)
+        if game.on_board(new_pos):
+            self.position = new_pos
+
+        game.state['energy'] -= 1
+        game.state['reward'] -= 0.01
+
+        food = game.get_agent_by_name("Food")
+        if self.position == food.position:
+            food.relocate(game)
+            game.state['energy'] += FOOD_ENERGY
+            game.state['score'] += 1
+            game.state['reward'] += 1.0
+
+        ax, ay = self.position
+        fx, fy = game.get_agent_by_name("Food").position
+        game.state['agent_x'] = ax
+        game.state['agent_y'] = ay
+        game.state['food_x'] = fx
+        game.state['food_y'] = fy
+
+        if game.state['energy'] <= 0:
+            game.end()
+
+
+class Food:
+    """The food item. Respawns at a random empty position when collected."""
+    name = "Food"
+    character = '*'
+    color = "yellow_on_black"
+    position = (0, 0)
+
+    def relocate(self, game):
+        bw, bh = game.board_size
+        forager = game.get_agent_by_name("Forager")
+        while True:
+            pos = (randint(0, bw - 1), randint(0, bh - 1))
+            if pos != forager.position:
+                self.position = pos
+                return
+
+
+def create_game():
+    """Return a fresh BabySnake game."""
+    forager = Forager()
+    food = Food()
+    bw = bh = BOARD_SIZE
+    game = Game(
+        [forager, food],
+        {
+            'score': 0,
+            'reward': 0.0,
+            'energy': START_ENERGY,
+            'agent_x': 0,
+            'agent_y': 0,
+            'food_x': 0,
+            'food_y': 0,
+        },
+        board_size=(bw, bh),
+        framerate=6,
+    )
+    forager.position = (randint(0, bw - 1), randint(0, bh - 1))
+    food.relocate(game)
+    ax, ay = forager.position
+    fx, fy = food.position
+    game.state['agent_x'] = ax
+    game.state['agent_y'] = ay
+    game.state['food_x'] = fx
+    game.state['food_y'] = fy
+    return game
+
+
+if __name__ == '__main__':
+    create_game().play()