myagent / agent.py
Kiss Mihaly
Upload folder using huggingface_hub
89801c9 verified
raw
history blame
1.22 kB
from stable_baselines3 import PPO
from sokoban_gym.wrappers.observation import ImageObservationWrapper
class Agent:
"""
A kötelező programként beadandó ágens leírása.
"""
def __init__(self, env) -> None:
"""
A konsztruktorban van lehetőség például a modell betöltésére
vagy a környezet wrapper-ekkel való kiterjesztésére.
"""
self.model = PPO.load("agent/Sokoban-A2C-CNN5")
# A környezetet kiterjeszthetjük wrapper-ek segítségével.
# Ha tanításkor modosítottuk a megfigyeléseket,
# akkor azt a módosítást kiértékeléskor is meg kell adnunk.
self.observation_wrapper = ImageObservationWrapper(env)
def act(self, observation):
"""
A megfigyelés alapján visszaadja a következő lépést.
Ez a függvény fogja megadni az ágens működését.
"""
# Ha tanításkor modosítottuk a megfigyeléseket,
# akkor azt a módosítást kiértékeléskor is meg kell adnunk.
extended_obsetvation = self.observation_wrapper.observation(observation)
return self.model.predict(extended_obsetvation, deterministic=True)