mirror of
https://github.com/wassname/DeepRL.git
synced 2026-09-09 11:13:47 +08:00
Major refactor
This commit is contained in:
Shangtong Zhang
+41
-69
@@ -10,9 +10,11 @@ import numpy as np
|
||||
import torch.multiprocessing as mp
|
||||
from task import *
|
||||
from network import *
|
||||
from bootstrap import *
|
||||
from worker import *
|
||||
import pickle
|
||||
import os
|
||||
import traceback
|
||||
import time
|
||||
|
||||
class AsyncAgent:
|
||||
def __init__(self,
|
||||
@@ -20,7 +22,7 @@ class AsyncAgent:
|
||||
network_fn,
|
||||
optimizer_fn,
|
||||
policy_fn,
|
||||
bootstrap,
|
||||
worker_fn,
|
||||
discount,
|
||||
step_limit,
|
||||
target_network_update_freq,
|
||||
@@ -33,13 +35,10 @@ class AsyncAgent:
|
||||
self.network_fn = network_fn
|
||||
self.learning_network = network_fn()
|
||||
self.learning_network.share_memory()
|
||||
if bootstrap != AdvantageActorCritic:
|
||||
self.target_network = network_fn()
|
||||
self.target_network.share_memory()
|
||||
self.target_network.load_state_dict(self.learning_network.state_dict())
|
||||
else:
|
||||
self.target_network = None
|
||||
self.bootstrap = bootstrap
|
||||
self.target_network = network_fn()
|
||||
self.target_network.share_memory()
|
||||
self.target_network.load_state_dict(self.learning_network.state_dict())
|
||||
self.worker_fn = worker_fn
|
||||
|
||||
self.optimizer_fn = optimizer_fn
|
||||
self.task_fn = task_fn
|
||||
@@ -66,86 +65,41 @@ class AsyncAgent:
|
||||
total_rewards = 0
|
||||
steps = 0
|
||||
network.reset(True)
|
||||
bootstrap = self.bootstrap(self)
|
||||
while not self.step_limit or steps < self.step_limit:
|
||||
action = np.argmax(bootstrap.process_state(network, state))
|
||||
action_value = network.predict(np.stack([state]))
|
||||
if self.worker_fn == AdvantageActorCritic:
|
||||
action_value = action_value[0]
|
||||
action = np.argmax(action_value.data.numpy().flatten())
|
||||
state, reward, terminal, _ = task.step(action)
|
||||
steps += 1
|
||||
total_rewards += reward
|
||||
if terminal:
|
||||
break
|
||||
bootstrap.reset()
|
||||
return total_rewards
|
||||
|
||||
def worker(self, id):
|
||||
optimizer = self.optimizer_fn(self.learning_network.parameters())
|
||||
worker_network = self.network_fn()
|
||||
worker_network.load_state_dict(self.learning_network.state_dict())
|
||||
|
||||
bootstrap = self.bootstrap(self)
|
||||
task = self.task_fn()
|
||||
policy = self.policy_fn()
|
||||
def train(self, id):
|
||||
worker = self.worker_fn(self)
|
||||
episode = 0
|
||||
episode_steps = 0
|
||||
episode_returns = [0]
|
||||
state = task.reset()
|
||||
pending_steps = 0
|
||||
rewards = []
|
||||
while True and not self.stop_signal.value:
|
||||
action = policy.sample(bootstrap.process_state(worker_network, state))
|
||||
next_state, reward, terminal, _ = task.step(action)
|
||||
bootstrap.process_interaction(action, reward, next_state)
|
||||
|
||||
episode_returns[-1] += reward
|
||||
episode_steps += 1
|
||||
if self.step_limit and episode_steps > self.step_limit:
|
||||
terminal = True
|
||||
with self.steps_lock:
|
||||
self.total_steps.value += 1
|
||||
pending_steps += 1
|
||||
|
||||
if terminal or pending_steps >= self.update_interval:
|
||||
loss = bootstrap.compute_loss(worker_network, terminal)
|
||||
pending_steps = 0
|
||||
worker_network.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm(worker_network.parameters(), 40)
|
||||
optimizer.zero_grad()
|
||||
for param, worker_param in zip(self.learning_network.parameters(), worker_network.parameters()):
|
||||
param._grad = worker_param.grad.clone().cpu()
|
||||
optimizer.step()
|
||||
worker_network.load_state_dict(self.learning_network.state_dict())
|
||||
worker_network.reset(terminal)
|
||||
|
||||
if terminal:
|
||||
state = task.reset()
|
||||
episode += 1
|
||||
if id == 0:
|
||||
self.logger.info('episode %d, return %f, avg return %f, episode steps %d, total steps %d' % (
|
||||
episode, episode_returns[-1], np.mean(episode_returns[-100:]), episode_steps, self.total_steps.value))
|
||||
episode_returns.append(0)
|
||||
episode_steps = 0
|
||||
else:
|
||||
state = next_state
|
||||
|
||||
if self.target_network and self.total_steps.value % self.target_network_update_freq == 0:
|
||||
self.target_network.load_state_dict(self.learning_network.state_dict())
|
||||
steps, reward = worker.episode()
|
||||
rewards.append(reward)
|
||||
self.logger.debug('worker %d, episode %d, return %f, avg return %f, episode steps %d, total steps %d' % (
|
||||
id, episode, rewards[-1], np.mean(rewards[-100:]), steps, self.total_steps.value))
|
||||
|
||||
def save(self, file_name):
|
||||
with open(file_name, 'wb') as f:
|
||||
pickle.dump(self.learning_network.state_dict(), f)
|
||||
|
||||
def run(self):
|
||||
os.environ['OMP_NUM_THREADS'] = '1'
|
||||
procs = [mp.Process(target=self.worker, args=(i, )) for i in range(self.n_workers)]
|
||||
for p in procs: p.start()
|
||||
def evaluate(self, id):
|
||||
test_rewards = []
|
||||
test_points = []
|
||||
test_network = self.network_fn()
|
||||
while True:
|
||||
steps = self.total_steps.value + 1
|
||||
steps = self.total_steps.value
|
||||
if steps % self.test_interval == 0:
|
||||
test_network.load_state_dict(self.learning_network.state_dict())
|
||||
self.save('data/%s%s-model-%s.bin' % (self.tag, self.bootstrap.__name__, self.task.name))
|
||||
self.save('data/%s%s-model-%s.bin' % (self.tag, self.worker_fn.__name__, self.task.name))
|
||||
rewards = np.zeros(self.test_repetitions)
|
||||
for i in range(self.test_repetitions):
|
||||
rewards[i] = self.deterministic_episode(self.task, test_network)
|
||||
@@ -154,10 +108,28 @@ class AsyncAgent:
|
||||
test_rewards.append(np.mean(rewards))
|
||||
test_points.append(steps)
|
||||
with open('data/%s%s-statistics-%s.bin' % (
|
||||
self.tag, self.bootstrap.__name__, self.task.name
|
||||
self.tag, self.worker_fn.__name__, self.task.name
|
||||
), 'wb') as f:
|
||||
pickle.dump([test_points, test_rewards], f)
|
||||
if np.mean(rewards) > self.task.success_threshold:
|
||||
self.stop_signal.value = True
|
||||
break
|
||||
|
||||
def run(self):
|
||||
os.environ['OMP_NUM_THREADS'] = '1'
|
||||
procs = [mp.Process(target=self.train, args=(i, )) for i in range(self.n_workers)]
|
||||
procs.append(mp.Process(target=self.evaluate, args=(self.n_workers, )))
|
||||
for p in procs: p.start()
|
||||
while True:
|
||||
time.sleep(1)
|
||||
for i, p in enumerate(procs):
|
||||
if not p.is_alive() and not self.stop_signal.value:
|
||||
self.logger.warning('Worker %d exited unexpectedly.' % i)
|
||||
p.terminate()
|
||||
procs[i] = mp.Process(target=self.train, args=(i, ))
|
||||
procs[i].start()
|
||||
self.logger.warning('Worker %d restarted.' % i)
|
||||
break
|
||||
if self.stop_signal.value:
|
||||
break
|
||||
for p in procs: p.join()
|
||||
|
||||
-139
@@ -1,139 +0,0 @@
|
||||
#######################################################################
|
||||
# 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 numpy as np
|
||||
import torch
|
||||
from torch.autograd import Variable
|
||||
|
||||
class OneStepSarsa:
|
||||
def __init__(self, agent):
|
||||
self.agent = agent
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.pending = []
|
||||
|
||||
def process_state(self, network, state):
|
||||
q = network.predict(np.stack([state]))
|
||||
self.pending.append([q])
|
||||
return q.data.numpy().flatten()
|
||||
|
||||
def process_interaction(self, action, reward, next_state):
|
||||
self.pending[-1].extend([action, reward, next_state])
|
||||
|
||||
def compute_loss(self, network, terminal):
|
||||
loss = 0
|
||||
valid_length = len(self.pending)
|
||||
if not terminal:
|
||||
valid_length -= 1
|
||||
for i in range(valid_length):
|
||||
q, action, reward, next_state = self.pending[i]
|
||||
q_next = self.agent.target_network.predict(np.stack([next_state])).data
|
||||
if i < len(self.pending) - 1:
|
||||
next_action = self.pending[i + 1][1]
|
||||
q_next = q_next.gather(1, torch.LongTensor([[next_action]]))
|
||||
else:
|
||||
q_next = torch.FloatTensor([[0]])
|
||||
q_next = self.agent.discount * q_next + reward
|
||||
q = q.gather(1, Variable(torch.LongTensor([[action]])))
|
||||
loss += 0.5 * (q - Variable(q_next)).pow(2)
|
||||
self.reset()
|
||||
return loss
|
||||
|
||||
|
||||
class OneStepQLearning:
|
||||
def __init__(self, agent):
|
||||
self.agent = agent
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.pending = []
|
||||
|
||||
def process_state(self, network, state):
|
||||
q = network.predict(np.stack([state]))
|
||||
self.pending.append([q])
|
||||
return q.data.numpy().flatten()
|
||||
|
||||
def process_interaction(self, action, reward, next_state):
|
||||
self.pending[-1].extend([action, reward, next_state])
|
||||
|
||||
def compute_loss(self, network, terminal):
|
||||
loss = 0
|
||||
for i in range(len(self.pending)):
|
||||
q, action, reward, next_state = self.pending[i]
|
||||
q_next, _ = self.agent.target_network.predict(np.stack([next_state])).data.max(1)
|
||||
if terminal and i == len(self.pending) - 1:
|
||||
q_next = torch.FloatTensor([[0]])
|
||||
q_next = self.agent.discount * q_next + reward
|
||||
q = q.gather(1, Variable(torch.LongTensor([[action]])))
|
||||
loss += 0.5 * (q - Variable(q_next)).pow(2)
|
||||
self.reset()
|
||||
return loss
|
||||
|
||||
class NStepQLearning:
|
||||
def __init__(self, agent):
|
||||
self.agent = agent
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.pending = []
|
||||
|
||||
def process_state(self, network, state):
|
||||
q = network.predict(np.stack([state]))
|
||||
self.pending.append([q])
|
||||
return q.data.numpy().flatten()
|
||||
|
||||
def process_interaction(self, action, reward, next_state):
|
||||
self.pending[-1].extend([action, reward])
|
||||
self.tailing_state = next_state
|
||||
|
||||
def compute_loss(self, network, terminal):
|
||||
loss = 0
|
||||
if terminal:
|
||||
R = torch.FloatTensor([[0]])
|
||||
else:
|
||||
R, _ = self.agent.target_network.predict(
|
||||
np.stack([self.tailing_state])).data.max(1)
|
||||
|
||||
for i in reversed(range(len(self.pending))):
|
||||
q, action, reward = self.pending[i]
|
||||
R = reward + self.agent.discount * R
|
||||
loss += 0.5 * (Variable(R) - q.gather(1, Variable(torch.LongTensor([[action]])))).pow(2)
|
||||
self.reset()
|
||||
return loss
|
||||
|
||||
class AdvantageActorCritic:
|
||||
def __init__(self, agent):
|
||||
self.agent = agent
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.pending = []
|
||||
|
||||
def process_state(self, network, state):
|
||||
prob, log_prob, value = network.predict(np.stack([state]))
|
||||
self.pending.append([prob, log_prob, value])
|
||||
return prob.data.numpy().flatten()
|
||||
|
||||
def process_interaction(self, action, reward, next_state):
|
||||
self.pending[-1].extend([action, reward])
|
||||
self.tailing_state = next_state
|
||||
|
||||
def compute_loss(self, network, terminal):
|
||||
loss = 0
|
||||
if terminal:
|
||||
R = torch.FloatTensor([[0]])
|
||||
else:
|
||||
R = network.critic(np.stack([self.tailing_state])).data
|
||||
for i in reversed(range(len(self.pending))):
|
||||
prob, log_prob, value, action, reward = self.pending[i]
|
||||
R = reward + self.agent.discount * R
|
||||
advantage = Variable(R) - value
|
||||
loss += 0.5 * advantage.pow(2)
|
||||
loss += -log_prob.gather(1, Variable(torch.LongTensor([[action]]))) * Variable(advantage.data)
|
||||
loss += 0.01 * torch.sum(torch.mul(prob, log_prob))
|
||||
self.reset()
|
||||
return loss
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from async_agent import *
|
||||
from dqn_agent import *
|
||||
import logging
|
||||
import traceback
|
||||
|
||||
def dqn_cart_pole():
|
||||
config = dict()
|
||||
@@ -29,9 +30,9 @@ def async_cart_pole():
|
||||
config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001)
|
||||
config['network_fn'] = lambda: FCNet([4, 50, 200, 2])
|
||||
config['policy_fn'] = lambda: GreedyPolicy(epsilon=0.5, final_step=5000, min_epsilon=0.1)
|
||||
config['bootstrap'] = OneStepQLearning
|
||||
# config['bootstrap'] = NStepQLearning
|
||||
# config['bootstrap'] = OneStepSarsa
|
||||
config['worker_fn'] = OneStepQLearning
|
||||
# config['worker_fn'] = NStepQLearning
|
||||
# config['worker_fn'] = OneStepSarsa
|
||||
config['discount'] = 0.99
|
||||
config['target_network_update_freq'] = 200
|
||||
config['step_limit'] = 0
|
||||
@@ -51,7 +52,7 @@ def a3c_cart_pole():
|
||||
config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001)
|
||||
config['network_fn'] = lambda: ActorCriticFCNet([4, 200, 2])
|
||||
config['policy_fn'] = SamplePolicy
|
||||
config['bootstrap'] = AdvantageActorCritic
|
||||
config['worker_fn'] = AdvantageActorCritic
|
||||
config['discount'] = 0.99
|
||||
config['target_network_update_freq'] = 200
|
||||
config['step_limit'] = 0
|
||||
@@ -96,13 +97,13 @@ def async_pixel_atari(name):
|
||||
config['optimizer_fn'] = lambda params: torch.optim.Adam(params, lr=0.0001)
|
||||
config['network_fn'] = lambda: OpenAIConvNet(history_length,
|
||||
n_actions)
|
||||
config['policy_fn'] = lambda: StochasticGreedyPolicy(epsilons=[0.5, 0.5, 0.5],
|
||||
config['policy_fn'] = lambda: StochasticGreedyPolicy(epsilons=[0.7, 0.7, 0.7],
|
||||
final_step=2000000,
|
||||
min_epsilons=[0.1, 0.01, 0.2],
|
||||
min_epsilons=[0.1, 0.01, 0.5],
|
||||
probs=[0.4, 0.3, 0.3])
|
||||
# config['bootstrap'] = OneStepQLearning
|
||||
# config['bootstrap'] = NStepQLearning
|
||||
config['bootstrap'] = OneStepSarsa
|
||||
# config['worker_fn'] = OneStepQLearning
|
||||
# config['worker_fn'] = NStepQLearning
|
||||
config['worker_fn'] = OneStepSarsa
|
||||
config['discount'] = 0.99
|
||||
config['target_network_update_freq'] = 10000
|
||||
config['step_limit'] = 10000
|
||||
@@ -113,6 +114,7 @@ def async_pixel_atari(name):
|
||||
config['history_length'] = history_length
|
||||
config['logger'] = gym.logger
|
||||
agent = AsyncAgent(**config)
|
||||
agent.tag = 'Centered-target-network-'
|
||||
agent.run()
|
||||
|
||||
def a3c_pixel_atari(name):
|
||||
@@ -125,7 +127,7 @@ def a3c_pixel_atari(name):
|
||||
n_actions,
|
||||
LSTM=False)
|
||||
config['policy_fn'] = SamplePolicy
|
||||
config['bootstrap'] = AdvantageActorCritic
|
||||
config['worker_fn'] = AdvantageActorCritic
|
||||
config['discount'] = 0.99
|
||||
config['target_network_update_freq'] = 0
|
||||
config['step_limit'] = 10000
|
||||
@@ -140,8 +142,8 @@ def a3c_pixel_atari(name):
|
||||
agent.run()
|
||||
|
||||
if __name__ == '__main__':
|
||||
# gym.logger.setLevel(logging.DEBUG)
|
||||
gym.logger.setLevel(logging.INFO)
|
||||
gym.logger.setLevel(logging.DEBUG)
|
||||
# gym.logger.setLevel(logging.INFO)
|
||||
|
||||
# dqn_cart_pole()
|
||||
# async_cart_pole()
|
||||
|
||||
@@ -0,0 +1,251 @@
|
||||
#######################################################################
|
||||
# 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 numpy as np
|
||||
import torch
|
||||
from torch.autograd import Variable
|
||||
import torch.nn as nn
|
||||
|
||||
class AdvantageActorCritic:
|
||||
def __init__(self, agent):
|
||||
self.agent = agent
|
||||
self.optimizer = agent.optimizer_fn(agent.learning_network.parameters())
|
||||
self.worker_network = agent.network_fn()
|
||||
self.worker_network.load_state_dict(agent.learning_network.state_dict())
|
||||
self.task = agent.task_fn()
|
||||
self.policy = agent.policy_fn()
|
||||
|
||||
def episode(self):
|
||||
state = self.task.reset()
|
||||
steps = 0
|
||||
total_reward = 0
|
||||
pending = []
|
||||
while True and not self.agent.stop_signal.value:
|
||||
prob, log_prob, value = self.worker_network.predict(np.stack([state]))
|
||||
action = self.policy.sample(prob.data.numpy().flatten())
|
||||
next_state, reward, terminal, _ = self.task.step(action)
|
||||
pending.append([prob, log_prob, value, action, reward])
|
||||
|
||||
steps += 1
|
||||
with self.agent.steps_lock:
|
||||
self.agent.total_steps.value += 1
|
||||
total_reward += reward
|
||||
|
||||
if terminal or len(pending) >= self.agent.update_interval:
|
||||
loss = 0
|
||||
if terminal:
|
||||
R = torch.FloatTensor([[0]])
|
||||
else:
|
||||
R = self.worker_network.critic(np.stack([next_state])).data
|
||||
for i in reversed(range(len(pending))):
|
||||
prob, log_prob, value, action, reward = pending[i]
|
||||
R = reward + self.agent.discount * R
|
||||
advantage = Variable(R) - value
|
||||
loss += 0.5 * advantage.pow(2)
|
||||
loss += -log_prob.gather(1, Variable(torch.LongTensor([[action]]))) * Variable(advantage.data)
|
||||
loss += 0.01 * torch.sum(torch.mul(prob, log_prob))
|
||||
|
||||
pending = []
|
||||
self.worker_network.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm(self.worker_network.parameters(), 40)
|
||||
self.optimizer.zero_grad()
|
||||
for param, worker_param in zip(
|
||||
self.agent.learning_network.parameters(), self.worker_network.parameters()):
|
||||
param._grad = worker_param.grad.clone()
|
||||
self.optimizer.step()
|
||||
self.worker_network.load_state_dict(self.agent.learning_network.state_dict())
|
||||
self.worker_network.reset(terminal)
|
||||
|
||||
if terminal:
|
||||
break
|
||||
else:
|
||||
state = next_state
|
||||
|
||||
return steps, total_reward
|
||||
|
||||
class NStepQLearning:
|
||||
def __init__(self, agent):
|
||||
self.agent = agent
|
||||
self.optimizer = agent.optimizer_fn(agent.learning_network.parameters())
|
||||
self.worker_network = agent.network_fn()
|
||||
self.worker_network.load_state_dict(agent.learning_network.state_dict())
|
||||
self.task = agent.task_fn()
|
||||
self.policy = agent.policy_fn()
|
||||
|
||||
def episode(self):
|
||||
state = self.task.reset()
|
||||
steps = 0
|
||||
total_reward = 0
|
||||
pending = []
|
||||
while True and not self.agent.stop_signal.value:
|
||||
q = self.worker_network.predict(np.stack([state]))
|
||||
action = self.policy.sample(q.data.numpy().flatten())
|
||||
next_state, reward, terminal, _ = self.task.step(action)
|
||||
pending.append([q, action, reward])
|
||||
|
||||
steps += 1
|
||||
with self.agent.steps_lock:
|
||||
self.agent.total_steps.value += 1
|
||||
total_reward += reward
|
||||
|
||||
if terminal or len(pending) >= self.agent.update_interval:
|
||||
loss = 0
|
||||
if terminal:
|
||||
R = torch.FloatTensor([[0]])
|
||||
else:
|
||||
R, _ = self.agent.target_network.predict(
|
||||
np.stack([next_state])).data.max(1)
|
||||
|
||||
for i in reversed(range(len(pending))):
|
||||
q, action, reward = pending[i]
|
||||
R = reward + self.agent.discount * R
|
||||
loss += 0.5 * (Variable(R) - q.gather(1, Variable(torch.LongTensor([[action]])))).pow(2)
|
||||
|
||||
pending = []
|
||||
self.worker_network.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm(self.worker_network.parameters(), 40)
|
||||
self.optimizer.zero_grad()
|
||||
for param, worker_param in zip(
|
||||
self.agent.learning_network.parameters(), self.worker_network.parameters()):
|
||||
param._grad = worker_param.grad.clone()
|
||||
self.optimizer.step()
|
||||
self.worker_network.load_state_dict(self.agent.learning_network.state_dict())
|
||||
self.worker_network.reset(terminal)
|
||||
|
||||
if terminal:
|
||||
break
|
||||
else:
|
||||
state = next_state
|
||||
|
||||
if self.agent.total_steps.value % self.agent.target_network_update_freq == 0:
|
||||
self.agent.target_network.load_state_dict(
|
||||
self.agent.learning_network.state_dict())
|
||||
|
||||
return steps, total_reward
|
||||
|
||||
class OneStepQLearning:
|
||||
def __init__(self, agent):
|
||||
self.agent = agent
|
||||
self.optimizer = agent.optimizer_fn(agent.learning_network.parameters())
|
||||
self.worker_network = agent.network_fn()
|
||||
self.worker_network.load_state_dict(agent.learning_network.state_dict())
|
||||
self.task = agent.task_fn()
|
||||
self.policy = agent.policy_fn()
|
||||
|
||||
def episode(self):
|
||||
state = self.task.reset()
|
||||
steps = 0
|
||||
total_reward = 0
|
||||
pending = []
|
||||
while True and not self.agent.stop_signal.value:
|
||||
q = self.worker_network.predict(np.stack([state]))
|
||||
action = self.policy.sample(q.data.numpy().flatten())
|
||||
next_state, reward, terminal, _ = self.task.step(action)
|
||||
pending.append([q, action, reward, next_state])
|
||||
|
||||
steps += 1
|
||||
with self.agent.steps_lock:
|
||||
self.agent.total_steps.value += 1
|
||||
total_reward += reward
|
||||
|
||||
if terminal or len(pending) >= self.agent.update_interval:
|
||||
loss = 0
|
||||
for i in range(len(pending)):
|
||||
q, action, reward, next_state = pending[i]
|
||||
q_next, _ = self.agent.target_network.predict(np.stack([next_state])).data.max(1)
|
||||
if terminal and i == len(pending) - 1:
|
||||
q_next = torch.FloatTensor([[0]])
|
||||
q_next = self.agent.discount * q_next + reward
|
||||
q = q.gather(1, Variable(torch.LongTensor([[action]])))
|
||||
loss += 0.5 * (q - Variable(q_next)).pow(2)
|
||||
|
||||
pending = []
|
||||
self.worker_network.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm(self.worker_network.parameters(), 40)
|
||||
self.optimizer.zero_grad()
|
||||
for param, worker_param in zip(
|
||||
self.agent.learning_network.parameters(), self.worker_network.parameters()):
|
||||
param._grad = worker_param.grad.clone()
|
||||
self.optimizer.step()
|
||||
self.worker_network.load_state_dict(self.agent.learning_network.state_dict())
|
||||
self.worker_network.reset(terminal)
|
||||
|
||||
if terminal:
|
||||
break
|
||||
else:
|
||||
state = next_state
|
||||
|
||||
if self.agent.total_steps.value % self.agent.target_network_update_freq == 0:
|
||||
self.agent.target_network.load_state_dict(
|
||||
self.agent.learning_network.state_dict())
|
||||
|
||||
return steps, total_reward
|
||||
|
||||
class OneStepSarsa:
|
||||
def __init__(self, agent):
|
||||
self.agent = agent
|
||||
self.optimizer = agent.optimizer_fn(agent.learning_network.parameters())
|
||||
self.worker_network = agent.network_fn()
|
||||
self.worker_network.load_state_dict(agent.learning_network.state_dict())
|
||||
self.task = agent.task_fn()
|
||||
self.policy = agent.policy_fn()
|
||||
|
||||
def episode(self):
|
||||
state = self.task.reset()
|
||||
q = self.worker_network.predict(np.stack([state]))
|
||||
action = self.policy.sample(q.data.numpy().flatten())
|
||||
steps = 0
|
||||
total_reward = 0
|
||||
pending = []
|
||||
while True and not self.agent.stop_signal.value:
|
||||
next_state, reward, terminal, _ = self.task.step(action)
|
||||
next_q = self.worker_network.predict(np.stack([next_state]))
|
||||
next_action = self.policy.sample(next_q.data.numpy().flatten())
|
||||
pending.append([q, action, reward, next_state, next_action])
|
||||
|
||||
steps += 1
|
||||
with self.agent.steps_lock:
|
||||
self.agent.total_steps.value += 1
|
||||
total_reward += reward
|
||||
|
||||
if terminal or len(pending) >= self.agent.update_interval:
|
||||
loss = 0
|
||||
for i in range(len(pending)):
|
||||
q, action, reward, next_state, next_action = pending[i]
|
||||
q_next = self.agent.target_network.predict(np.stack([next_state])).data
|
||||
if terminal and i == len(pending) - 1:
|
||||
q_next = torch.FloatTensor([[0]])
|
||||
else:
|
||||
q_next = q_next.gather(1, torch.LongTensor([[next_action]]))
|
||||
q_next = self.agent.discount * q_next + reward
|
||||
q = q.gather(1, Variable(torch.LongTensor([[action]])))
|
||||
loss += 0.5 * (q - Variable(q_next)).pow(2)
|
||||
|
||||
pending = []
|
||||
self.worker_network.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm(self.worker_network.parameters(), 40)
|
||||
self.optimizer.zero_grad()
|
||||
for param, worker_param in zip(
|
||||
self.agent.learning_network.parameters(), self.worker_network.parameters()):
|
||||
param._grad = worker_param.grad.clone()
|
||||
self.optimizer.step()
|
||||
self.worker_network.load_state_dict(self.agent.learning_network.state_dict())
|
||||
self.worker_network.reset(terminal)
|
||||
|
||||
if terminal:
|
||||
break
|
||||
else:
|
||||
q = next_q
|
||||
action = next_action
|
||||
|
||||
if self.agent.total_steps.value % self.agent.target_network_update_freq == 0:
|
||||
self.agent.target_network.load_state_dict(
|
||||
self.agent.learning_network.state_dict())
|
||||
|
||||
return steps, total_reward
|
||||
Reference in New Issue
Block a user