118 lines
3.1 KiB
Python
118 lines
3.1 KiB
Python
"""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()
|