Major reversion

This commit is contained in:
Shangtong Zhang
2017-07-30 13:24:04 -06:00
parent ce504e2d0f
commit bbfcbf76f4
7 changed files with 131 additions and 110 deletions
+9 -7
View File
@@ -9,13 +9,15 @@ from torch.autograd import Variable
import torch.nn as nn
class OneStepQLearning:
def __init__(self, config):
def __init__(self, config, learning_network, target_network):
self.config = config
self.optimizer = config.optimizer_fn(config.learning_network.parameters())
self.optimizer = config.optimizer_fn(learning_network.parameters())
self.worker_network = config.network_fn()
self.worker_network.load_state_dict(config.learning_network.state_dict())
self.worker_network.load_state_dict(learning_network.state_dict())
self.task = config.task_fn()
self.policy = config.policy_fn()
self.learning_network = learning_network
self.target_network = target_network
def episode(self, deterministic=False):
config = self.config
@@ -46,7 +48,7 @@ class OneStepQLearning:
loss = 0
for i in range(len(pending)):
q, action, reward, next_state = pending[i]
q_next, _ = config.target_network.predict(np.stack([next_state])).data.max(1)
q_next, _ = self.target_network.predict(np.stack([next_state])).data.max(1)
if terminal and i == len(pending) - 1:
q_next = torch.FloatTensor([[0]])
q_next = config.discount * q_next + reward
@@ -59,12 +61,12 @@ class OneStepQLearning:
loss.backward()
nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
for param, worker_param in zip(
config.learning_network.parameters(), self.worker_network.parameters()):
self.learning_network.parameters(), self.worker_network.parameters()):
if param.grad is not None:
break
param._grad = worker_param.grad
self.optimizer.step()
self.worker_network.load_state_dict(config.learning_network.state_dict())
self.worker_network.load_state_dict(self.learning_network.state_dict())
self.worker_network.reset(terminal)
if terminal:
@@ -72,6 +74,6 @@ class OneStepQLearning:
state = next_state
if config.total_steps.value % config.target_network_update_freq == 0:
config.target_network.load_state_dict(config.learning_network.state_dict())
self.target_network.load_state_dict(self.learning_network.state_dict())
return steps, total_reward