mirror of
https://github.com/wassname/DeepRL.git
synced 2026-08-21 11:09:46 +08:00
114 lines
5.1 KiB
Python
114 lines
5.1 KiB
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 numpy as np
|
|
import torch
|
|
from torch.autograd import Variable
|
|
import torch.nn as nn
|
|
from utils import *
|
|
|
|
class ContinuousAdvantageActorCritic:
|
|
def __init__(self, config, learning_network, extra):
|
|
self.config = config
|
|
self.actor_opt = config.actor_optimizer_fn(learning_network.actor.parameters())
|
|
self.critic_opt = config.critic_optimizer_fn(learning_network.critic.parameters())
|
|
self.worker_network = config.network_fn()
|
|
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.counter = 0
|
|
|
|
self.shared_state_normalizer = extra[0]
|
|
self.state_normalizer = StaticNormalizer(self.task.state_dim)
|
|
self.shared_reward_normalizer = extra[1]
|
|
self.reward_normalizer = StaticNormalizer(1)
|
|
|
|
def episode(self, deterministic=False):
|
|
config = self.config
|
|
self.state_normalizer.offline_stats.load(self.shared_state_normalizer.offline_stats)
|
|
self.reward_normalizer.offline_stats.load(self.shared_reward_normalizer.offline_stats)
|
|
state = self.task.reset()
|
|
state = self.state_normalizer(state)
|
|
steps = 0
|
|
total_reward = 0
|
|
pending = []
|
|
while not config.stop_signal.value:
|
|
mean, std, log_std = self.worker_network.actor.predict(np.stack([state]))
|
|
value = self.worker_network.critic.predict(np.stack([state]))
|
|
action = self.policy.sample(mean.data.numpy().flatten(),
|
|
std.data.numpy().flatten(),
|
|
False)
|
|
action = self.config.action_shift_fn(action)
|
|
next_state, reward, terminal, _ = self.task.step(action)
|
|
terminal = (terminal or (config.max_episode_length and steps >= config.max_episode_length))
|
|
next_state = self.state_normalizer(next_state)
|
|
|
|
steps += 1
|
|
total_reward += reward
|
|
reward = self.reward_normalizer(reward)
|
|
|
|
if deterministic:
|
|
if terminal:
|
|
break
|
|
state = next_state
|
|
continue
|
|
|
|
pending.append([mean, std, log_std, value, action, reward])
|
|
with config.steps_lock:
|
|
config.total_steps.value += 1
|
|
|
|
if terminal or len(pending) >= config.update_interval:
|
|
critic_loss = 0
|
|
actor_loss = 0
|
|
if terminal:
|
|
R = torch.FloatTensor([[0]])
|
|
else:
|
|
R = self.worker_network.critic(np.stack([next_state])).data
|
|
GAE = torch.FloatTensor([[0]])
|
|
for i in reversed(range(len(pending))):
|
|
mean, std, log_std, value, action, reward = pending[i]
|
|
if i == len(pending) - 1:
|
|
delta = reward + config.discount * R - value.data
|
|
else:
|
|
delta = reward + pending[i + 1][3].data - value.data
|
|
GAE = config.discount * config.gae_tau * GAE + delta
|
|
|
|
action = Variable(torch.FloatTensor([action]))
|
|
log_density = self.worker_network.actor.log_density(action, mean, log_std, std)
|
|
actor_loss += -torch.sum(log_density) * Variable(GAE)
|
|
if config.entropy_weight:
|
|
actor_loss += -config.entropy_weight * self.worker_network.actor.entropy(std)
|
|
|
|
R = reward + config.discount * R
|
|
critic_loss += 0.5 * (Variable(R) - value).pow(2)
|
|
|
|
pending = []
|
|
self.worker_network.zero_grad()
|
|
self.actor_opt.zero_grad()
|
|
self.critic_opt.zero_grad()
|
|
actor_loss.backward()
|
|
critic_loss.backward()
|
|
nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
|
|
for param, worker_param in zip(
|
|
self.learning_network.parameters(), self.worker_network.parameters()):
|
|
if param.grad is not None:
|
|
break
|
|
param._grad = worker_param.grad
|
|
self.actor_opt.step()
|
|
self.critic_opt.step()
|
|
self.worker_network.load_state_dict(self.learning_network.state_dict())
|
|
|
|
if terminal:
|
|
break
|
|
state = next_state
|
|
|
|
self.shared_state_normalizer.offline_stats.merge(self.state_normalizer.online_stats)
|
|
self.state_normalizer.online_stats.zero()
|
|
|
|
self.shared_reward_normalizer.offline_stats.merge(self.reward_normalizer.online_stats)
|
|
self.reward_normalizer.online_stats.zero()
|
|
|
|
return steps, total_reward |