92 lines
3.6 KiB
Python
92 lines
3.6 KiB
Python
import numpy as np
|
||
from retro.views.headless import HeadlessView
|
||
from retro_gamer.observation import egocentric_board, encode_board, encode_state
|
||
|
||
HEAD_CHARS = frozenset({">", "<", "^", "v"})
|
||
CHARACTER_SET = ["@", "*", ">", "<", "^", "v"]
|
||
CHARACTER_SET_NORMALIZED = ["@", "*", "H"]
|
||
RADIUS = 8
|
||
|
||
DIRECTIONS = [(1, 0), (0, -1), (-1, 0), (0, 1)] # RIGHT, UP, LEFT, DOWN
|
||
RADIUS_WIDE = 8 # 17×17 egocentric window
|
||
RADIUS_NARROW = 3 # 7×7 egocentric window
|
||
|
||
|
||
def direction_observation(game):
|
||
"""Normalized 3-char board plus apple_dx, apple_dy, and a one-hot direction.
|
||
|
||
Returns board in channel-first (C, H, W) order for the CNN, followed by
|
||
apple_dx, apple_dy, and 4 one-hot bits encoding the snake's heading.
|
||
The board collapses all four head characters to 'H', keeping 3 channels
|
||
instead of 6 while the one-hot direction restores that information cheaply.
|
||
"""
|
||
view = HeadlessView()
|
||
view.on_game_start(game)
|
||
view.render(game)
|
||
normalized = [
|
||
["H" if c in HEAD_CHARS else c for c in row]
|
||
for row in view.board_characters
|
||
]
|
||
board_vec = encode_board(normalized, CHARACTER_SET_NORMALIZED).transpose(2, 0, 1).flatten()
|
||
extras = encode_state(game.state, ["apple_dx", "apple_dy"])
|
||
head = game.get_agent_by_name("Snake head")
|
||
direction_onehot = np.array(
|
||
[1.0 if head.direction == d else 0.0 for d in DIRECTIONS],
|
||
dtype=np.float32,
|
||
)
|
||
return np.concatenate([board_vec, extras, direction_onehot])
|
||
|
||
|
||
def normalized_observation(game):
|
||
"""Full board with all head chars collapsed to 'H', plus apple_dx and apple_dy.
|
||
|
||
Returns board in channel-first (C, H, W) order for the CNN, followed by extras.
|
||
Character set: ['@' apple, '*' body, 'H' head] — 3 channels instead of 6.
|
||
"""
|
||
view = HeadlessView()
|
||
view.on_game_start(game)
|
||
view.render(game)
|
||
normalized = [
|
||
["H" if c in HEAD_CHARS else c for c in row]
|
||
for row in view.board_characters
|
||
]
|
||
board_vec = encode_board(normalized, CHARACTER_SET_NORMALIZED).transpose(2, 0, 1).flatten()
|
||
extras = encode_state(game.state, ["apple_dx", "apple_dy"])
|
||
return np.concatenate([board_vec, extras])
|
||
|
||
|
||
def egocentric_cnn_observation(game):
|
||
"""17×17 egocentric window in channel-first (CHW) format for CNN, plus apple_dx/dy."""
|
||
view = HeadlessView()
|
||
view.on_game_start(game)
|
||
view.render(game)
|
||
head = game.get_agent_by_name("Snake head")
|
||
cropped = egocentric_board(view.board_characters, head.position, RADIUS_WIDE)
|
||
board_vec = encode_board(cropped, CHARACTER_SET).transpose(2, 0, 1).flatten()
|
||
extras = encode_state(game.state, ["apple_dx", "apple_dy"])
|
||
return np.concatenate([board_vec, extras])
|
||
|
||
|
||
def narrow_egocentric_observation(game):
|
||
"""7×7 egocentric window (flat, for MLP), plus apple_dx and apple_dy."""
|
||
view = HeadlessView()
|
||
view.on_game_start(game)
|
||
view.render(game)
|
||
head = game.get_agent_by_name("Snake head")
|
||
cropped = egocentric_board(view.board_characters, head.position, RADIUS_NARROW)
|
||
board_vec = encode_board(cropped, CHARACTER_SET).flatten()
|
||
extras = encode_state(game.state, ["apple_dx", "apple_dy"])
|
||
return np.concatenate([board_vec, extras])
|
||
|
||
|
||
def egocentric_observation(game):
|
||
"""17×17 window centered on the snake's head, plus apple_dx and apple_dy."""
|
||
view = HeadlessView()
|
||
view.on_game_start(game)
|
||
view.render(game)
|
||
head = game.get_agent_by_name("Snake head")
|
||
cropped = egocentric_board(view.board_characters, head.position, RADIUS)
|
||
board_vec = encode_board(cropped, CHARACTER_SET).flatten()
|
||
extras = encode_state(game.state, ["apple_dx", "apple_dy"])
|
||
return np.concatenate([board_vec, extras])
|