mirror of
https://github.com/wassname/DeepRL.git
synced 2026-09-11 11:53:01 +08:00
Major reversion
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user