mirror of
https://github.com/wassname/DeepRL.git
synced 2026-08-21 11:09:46 +08:00
73 lines
2.8 KiB
Python
73 lines
2.8 KiB
Python
#######################################################################
|
|
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
|
|
# Permission given to modify the code as long as you keep this #
|
|
# declaration at the top #
|
|
#######################################################################
|
|
|
|
import torch
|
|
import numpy as np
|
|
from ..utils import *
|
|
|
|
class BaseAgent:
|
|
def __init__(self, config):
|
|
self.config = config
|
|
self.evaluation_env = self.config.evaluation_env
|
|
if self.evaluation_env is not None:
|
|
self.evaluation_state = self.evaluation_env.reset()
|
|
self.evaluation_return = 0
|
|
|
|
def close(self):
|
|
if hasattr(self.task, 'close'):
|
|
self.task.close()
|
|
if hasattr(self.evaluation_env, 'close'):
|
|
self.evaluation_env.close()
|
|
|
|
def save(self, filename):
|
|
torch.save(self.network.state_dict(), filename)
|
|
|
|
def load(self, filename):
|
|
state_dict = torch.load(filename, map_location=lambda storage, loc: storage)
|
|
self.network.load_state_dict(state_dict)
|
|
|
|
def evaluation_action(self, state):
|
|
self.config.state_normalizer.set_read_only()
|
|
state = self.config.state_normalizer(np.stack([state]))
|
|
action = self.network.predict(state, to_numpy=True)
|
|
self.config.state_normalizer.unset_read_only()
|
|
return np.argmax(action.flatten())
|
|
|
|
def deterministic_episode(self):
|
|
env = self.config.evaluation_env
|
|
state = env.reset()
|
|
total_rewards = 0
|
|
while True:
|
|
action = self.evaluation_action(state)
|
|
state, reward, done, _ = env.step(action)
|
|
total_rewards += reward
|
|
if done:
|
|
break
|
|
return total_rewards
|
|
|
|
def evaluation_episodes(self):
|
|
interval = self.config.evaluation_episodes_interval
|
|
if not interval or self.total_steps % interval:
|
|
return
|
|
rewards = []
|
|
for ep in range(self.config.evaluation_episodes):
|
|
rewards.append(self.deterministic_episode())
|
|
self.config.logger.info('evaluation episode return: %f(%f)' % (
|
|
np.mean(rewards), np.std(rewards) / np.sqrt(len(rewards))))
|
|
|
|
def evaluate(self, steps=1):
|
|
config = self.config
|
|
if config.evaluation_env is None or self.config.evaluation_episodes_interval:
|
|
return
|
|
for _ in range(steps):
|
|
action = self.evaluation_action(self.evaluation_state)
|
|
self.evaluation_state, reward, done, _ = self.evaluation_env.step(action)
|
|
self.evaluation_return += reward
|
|
if done:
|
|
self.evaluation_state = self.evaluation_env.reset()
|
|
self.config.logger.info('evaluation episode return: %f' % (self.evaluation_return))
|
|
self.evaluation_return = 0
|