mirror of
https://github.com/wassname/DeepRL.git
synced 2026-08-21 11:09:46 +08:00
22 lines
772 B
Python
22 lines
772 B
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
|
|
|
|
class BaseAgent:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def close(self):
|
|
if hasattr(self.task, 'close'):
|
|
self.task.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) |