def load_weights(self, best=True, load_buffer=True):
'''
Load the model weights and replay buffer from self.save_dir
Args:
best (bool): If True, save from the weights file with the best mean episode reward
load_buffer (bool): If True, load the replay buffer from the pickled file
'''
if best:
fname = "best.pth"
else:
fname = "model_weights.pth"
checkpoint_path = os.path.join(self.save_dir, fname)
if os.path.isfile(checkpoint_path):
if load_buffer:
self.replay_buffer.load(os.path.join(self.save_dir, "replay_buffer.pickle"))
key = 'cuda' if torch.cuda.is_available() else 'cpu'
checkpoint = torch.load(checkpoint_path, map_location=key)
self.ac.load_state_dict(sanitise_state_dict(checkpoint['ac'], self.ngpu>1))
self.ac_targ.load_state_dict(sanitise_state_dict(checkpoint['ac_target'], self.ngpu>1))
self.pi_optimizer.load_state_dict(sanitise_state_dict(checkpoint['pi_optimizer'], self.ngpu>1))
self.q_optimizer.load_state_dict(sanitise_state_dict(checkpoint['q_optimizer'], self.ngpu>1))
env_path = os.path.join(self.save_dir, "env.json")
if os.path.isfile(env_path):
self.env = self.env.load(env_path)
print("Environment loaded")
print('checkpoint loaded at {}'.format(checkpoint_path))
else:
raise OSError("Checkpoint file not found.")
def load_weights(self, best=True, fname=None):
'''
Load the model weights and replay buffer from self.save_dir
Args:
best (bool): If True, save from the weights file with the best mean episode reward
load_buffer (bool): If True, load the replay buffer from the pickled file
'''
if fname is not None:
_fname = fname
elif best:
_fname = "best.pth"
else:
_fname = "model_weights.pth"
checkpoint_path = os.path.join(self.save_dir, _fname)
if os.path.isfile(checkpoint_path):
checkpoint = torch.load(checkpoint_path, map_location=self.device)
self.network.load_state_dict(sanitise_state_dict(checkpoint['oc']))
self.target_network.load_state_dict(
sanitise_state_dict(checkpoint['oc_target']))
env_path = os.path.join(self.save_dir, "env.json")
if os.path.isfile(env_path):
self.env = self.env.load(env_path)
print("Environment loaded")
print('checkpoint loaded at {}'.format(checkpoint_path))
else:
raise OSError("Checkpoint file not found.")