<reponame>aarav18/snake-rl
from gym import Env
from gym.spaces import Discrete, Box
import numpy as np
import random
from snake import Game

class SnakeEnv(Env):
    def __init__(self):
        self.action_space = Discrete(3) # 0 = turn left, 1 = do nothing, 2 = turn right
        self.state = [0, 0, 1, 0]
        self.game = Game()
        self.reward = 0
        self.done = False

    def step(self, action):
        
        offset = (action - 1)
        translated_action = offset + self.game.snake.direction
        if translated_action < 0:
            translated_action = 3
        if translated_action > 3:
            translated_action = 0

        self.reward, self.done = self.game.run(1, translated_action)

        diff = (self.game.food.position[0] - self.game.snake.snake[0][0], self.game.food.position[1] - self.game.snake.snake[0][1])
        
        self.state[0] = int(diff[0] < 0)
        self.state[2] = int(diff[0] > 1)

        self.state[1] = int(diff[1] < 0)
        self.state[3] = int(diff[1] > 0)

        return self.state, self.reward, self.done, {}

    def render(self):
        self.game.render()

    def reset(self):
        self.game.reset()

if __name__ == "__main__":

    env = SnakeEnv()

    episodes = 2

    for episode in range(1, episodes+1):
        state = env.reset()
        done = False

        while not done:
            env.render()
            action = env.action_space.sample()
            n_state, reward, done, info = env.step(action)
        print("Episode: {}, Score: {}".format(episode, reward))
from gym import Env
from gym.spaces import Discrete, Box
import numpy as np
import random
from snake import Game

class SnakeEnv(Env):
    def __init__(self):
        self.action_space = Discrete(3) # 0 = turn left, 1 = do nothing, 2 = turn right
        self.state = [0, 0, 1, 0]
        self.game = Game()
        self.reward = 0
        self.done = False

    def step(self, action):
        
        offset = (action - 1)
        translated_action = offset + self.game.snake.direction
        if translated_action < 0:
            translated_action = 3
        if translated_action > 3:
            translated_action = 0

        self.reward, self.done = self.game.run(1, translated_action)

        diff = (self.game.food.position[0] - self.game.snake.snake[0][0], self.game.food.position[1] - self.game.snake.snake[0][1])
        
        self.state[0] = int(diff[0] < 0)
        self.state[2] = int(diff[0] > 1)

        self.state[1] = int(diff[1] < 0)
        self.state[3] = int(diff[1] > 0)

        return self.state, self.reward, self.done, {}

    def render(self):
        self.game.render()

    def reset(self):
        self.game.reset()

if __name__ == "__main__":

    env = SnakeEnv()

    episodes = 2

    for episode in range(1, episodes+1):
        state = env.reset()
        done = False

        while not done:
            env.render()
            action = env.action_space.sample()
            n_state, reward, done, info = env.step(action)
        print("Episode: {}, Score: {}".format(episode, reward))
# 0 = turn left, 1 = do nothing, 2 = turn right