WRITER.add_scalar('training/generator/loss', average_generator_loss,
epoch)
WRITER.add_scalar('training/critic/loss', average_critic_loss, epoch)
WRITER.add_scalar('training/critic/real-performance',
average_critic_real_performance, epoch)
WRITER.add_scalar('training/critic/generated-performance',
average_critic_generated_performance, epoch)
WRITER.add_scalar('training/discernability-score',
discernability_score, epoch)
WRITER.add_scalar('training/epoch-duration', time_elapsed, epoch)
# Save the model parameters at a specified interval.
if (epoch > 0 and (epoch % args.model_save_frequency == 0
or epoch == args.num_epochs - 1)):
# Create a backup of tensorboard data each time model is saved.
shutil.copytree(LIVE_TENSORBOARD_DIR,
f'{TENSORBOARD_DIR}/{epoch:03d}')
save_critic_model_path = f'{args.save_model_dir}/critic_{EXPERIMENT_ID}-{epoch}.pth'
print(f'\nSaving critic model as "{save_critic_model_path}"...')
torch.save(critic_model.state_dict(), save_critic_model_path)
save_generator_model_path = f'{args.save_model_dir}/generator_{EXPERIMENT_ID}-{epoch}.pth'
print(
f'Saving generator model as "{save_generator_model_path}"...\n'
)
torch.save(generator_model.state_dict(), save_generator_model_path)
print('Finished training!')
class Agent():
"""Interacts with and learns from the environment."""
def __init__(self, state_size: int, action_size: int, seed: int,
n_agent: int):
"""Initialize an Agent object.
Params
======
state_size (int): dimension of each state
action_size (int): dimension of each action
seed (int): random seed
"""
self.state_size = state_size
self.action_size = action_size
self.n_agent = n_agent
self.seed = random.seed(seed)
self.global_step = 0
self.update_step = 0
# Initialize actor and critic local and target networks
self.actor = Actor(state_size,
action_size,
seed,
ACTOR_NETWORK_LINEAR_SIZES,
batch_normalization=ACTOR_BATCH_NORM).to(device)
self.actor_target = Actor(
state_size,
action_size,
seed,
ACTOR_NETWORK_LINEAR_SIZES,
batch_normalization=ACTOR_BATCH_NORM).to(device)
self.critic = Critic(state_size,
action_size,
seed,
CRITIC_NETWORK_LINEAR_SIZES,
batch_normalization=CRITIC_BATCH_NORM).to(device)
self.critic_second = Critic(
state_size,
action_size,
seed,
CRITIC_SECOND_NETWORK_LINEAR_SIZES,
batch_normalization=CRITIC_BATCH_NORM).to(device)
self.critic_second_target = Critic(
state_size,
action_size,
seed,
CRITIC_SECOND_NETWORK_LINEAR_SIZES,
batch_normalization=CRITIC_BATCH_NORM).to(device)
self.critic_target = Critic(
state_size,
action_size,
seed,
CRITIC_NETWORK_LINEAR_SIZES,
batch_normalization=CRITIC_BATCH_NORM).to(device)
self.actor_optimizer = optim.Adam(self.actor.parameters(),
lr=ACTOR_LEARNING_RATE)
self.critic_optimizer = optim.Adam(self.critic.parameters(),
lr=CRITIC_LEARNING_RATE)
self.critic_second_optimizer = optim.Adam(
self.critic_second.parameters(), lr=CRITIC_LEARNING_RATE)
# Replay memory
self.memory = ReplayBuffer(action_size, BUFFER_SIZE, BATCH_SIZE, seed)
# Initialize time step (for updating every UPDATE_EVERY steps)
self.t_step = [0] * n_agent
self.noise = OUNoise(action_size, seed, decay_period=50)
# Copy parameters from local network to target network
for target_param, param in zip(self.actor_target.parameters(),
self.actor.parameters()):
target_param.data.copy_(param.data)
for target_param, param in zip(self.critic_target.parameters(),
self.critic.parameters()):
target_param.data.copy_(param.data)
for target_param, param in zip(self.critic_second_target.parameters(),
self.critic_second.parameters()):
target_param.data.copy_(param.data)
def step(self, state: np.array, action, reward, next_state, done, i_agent):
# Save experience in replay memory
self.memory.add(state, action, reward, next_state, done)
# Learn every UPDATE_EVERY time steps.
self.t_step[i_agent] = (self.t_step[i_agent] + 1) % UPDATE_EVERY
# Learn, if enough samples are available in memory every UPDATE_EVERY
if len(self.memory) > BATCH_SIZE and (not any(self.t_step)):
for _ in range(LEARN_TIMES):
experiences = self.memory.sample()
self.learn(experiences, GAMMA)
def noise_reset(self):
self.noise.reset()
def save_model(self, checkpoint_path: str = "./checkpoints/"):
torch.save(self.actor.state_dict(), f"{checkpoint_path}/actor.pt")
torch.save(self.critic.state_dict(), f"{checkpoint_path}/critic.pt")
def load_model(self, checkpoint_path: str = "./checkpoints/checkpoint.pt"):
self.actor.load_state_dict(torch.load(f"{checkpoint_path}/actor.pt"))
self.critic.load_state_dict(torch.load(f"{checkpoint_path}/critic.pt"))
def act(self, states: np.array, step: int):
"""Returns actions for given state as per current policy.
Params
======
state (array_like): current state
"""
self.global_step += 1
if self.global_step < WARM_UP_STEPS:
action_values = np.random.rand(self.n_agent, self.action_size)
return action_values
states = torch.from_numpy(states).float().to(device)
self.actor.eval()
with torch.no_grad():
action_values = self.actor(states).cpu().data.numpy()
self.actor.train()
action_values = [
self.noise.get_action(action, t=step) for action in action_values
]
return action_values
def learn(self, experiences: tuple, gamma=GAMMA):
"""Update value parameters using given batch of experience tuples.
Params
======
experiences (Tuple[torch.Variable]): tuple of (s, a, r, s', done) tuples
gamma (float): discount factor
"""
self.update_step += 1
states, actions, rewards, next_states, dones = experiences
# Critic loss
mask = torch.tensor(1 - dones).detach().to(device)
Q_values = self.critic(states, actions)
Q_values_second = self.critic_second(states, actions)
next_actions = self.actor_target(next_states)
next_Q = torch.min(
self.critic_target(next_states, next_actions.detach()),
self.critic_second_target(next_states, next_actions.detach()))
Q_prime = rewards + gamma * next_Q * mask
if self.update_step % CRITIC_UPDATE_EVERY == 0:
critic_loss = F.mse_loss(Q_values, Q_prime.detach())
critic_second_loss = F.mse_loss(Q_values_second, Q_prime.detach())
# Update first critic network
self.critic_optimizer.zero_grad()
critic_loss.backward()
if CRITIC_GRADIENT_CLIPPING_VALUE:
torch.nn.utils.clip_grad_norm_(self.critic.parameters(),
CRITIC_GRADIENT_CLIPPING_VALUE)
self.critic_optimizer.step()
# Update second critic network
self.critic_second_optimizer.zero_grad()
critic_second_loss.backward()
if CRITIC_GRADIENT_CLIPPING_VALUE:
torch.nn.utils.clip_grad_norm_(self.critic_second.parameters(),
CRITIC_GRADIENT_CLIPPING_VALUE)
self.critic_second_optimizer.step()
# Actor loss
if self.update_step % POLICY_UPDATE_EVERY == 0:
policy_loss = -self.critic(states, self.actor(states)).mean()
# Update actor network
self.actor_optimizer.zero_grad()
policy_loss.backward()
if ACTOR_GRADIENT_CLIPPING_VALUE:
torch.nn.utils.clip_grad_norm_(self.actor.parameters(),
ACTOR_GRADIENT_CLIPPING_VALUE)
self.actor_optimizer.step()
self.actor_soft_update()
self.critic_soft_update()
self.critic_second_soft_update()
def actor_soft_update(self, tau: float = TAU):
"""Soft update for actor target network
Args:
tau (float, optional). Defaults to TAU.
"""
for target_param, param in zip(self.actor_target.parameters(),
self.actor.parameters()):
target_param.data.copy_(tau * param.data +
(1.0 - tau) * target_param.data)
def critic_soft_update(self, tau: float = TAU):
"""Soft update for critic target network
Args:
tau (float, optional). Defaults to TAU.
"""
for target_param, param in zip(self.critic_target.parameters(),
self.critic.parameters()):
target_param.detach_()
target_param.data.copy_(tau * param.data +
(1.0 - tau) * target_param.data)
def critic_second_soft_update(self, tau: float = TAU):
"""Soft update for critic target network
Args:
tau (float, optional). Defaults to TAU.
"""
for target_param, param in zip(self.critic_second_target.parameters(),
self.critic_second.parameters()):
target_param.detach_()
target_param.data.copy_(tau * param.data +
(1.0 - tau) * target_param.data)
class DDPGAgent(object):
"""
General class for DDPG agents (policy, critic, target policy, target
critic, exploration noise)
"""
def __init__(self, num_in_pol, num_out_pol, num_in_critic, hidden_dim_actor=120,
hidden_dim_critic=64,lr_actor=0.01,lr_critic=0.01,batch_size=64,
max_episode_len=100,tau=0.02,gamma = 0.99,agent_name='one', discrete_action=False):
"""
Inputs:
num_in_pol (int): number of dimensions for policy input
num_out_pol (int): number of dimensions for policy output
num_in_critic (int): number of dimensions for critic input
"""
self.policy = Actor(num_in_pol, num_out_pol,
hidden_dim=hidden_dim_actor,
discrete_action=discrete_action)
self.critic = Critic(num_in_pol, 1,num_out_pol,
hidden_dim=hidden_dim_critic)
self.target_policy = Actor(num_in_pol, num_out_pol,
hidden_dim=hidden_dim_actor,
discrete_action=discrete_action)
self.target_critic = Critic(num_in_pol, 1,num_out_pol,
hidden_dim=hidden_dim_critic)
hard_update(self.target_policy, self.policy)
hard_update(self.target_critic, self.critic)
self.policy_optimizer = Adam(self.policy.parameters(), lr=lr_actor)
self.critic_optimizer = Adam(self.critic.parameters(), lr=lr_critic,weight_decay=0)
self.policy = self.policy.float()
self.critic = self.critic.float()
self.target_policy = self.target_policy.float()
self.target_critic = self.target_critic.float()
self.agent_name = agent_name
self.gamma = gamma
self.tau = tau
self.batch_size = batch_size
#self.replay_buffer = ReplayBuffer(1e7)
self.replay_buffer = ReplayBufferOption(500000,self.batch_size,12)
self.max_replay_buffer_len = batch_size * max_episode_len
self.replay_sample_index = None
self.niter = 0
self.eps = 5.0
self.eps_decay = 1/(250*5)
self.exploration = OUNoise(num_out_pol)
self.discrete_action = discrete_action
self.num_history = 2
self.states = []
self.actions = []
self.rewards = []
self.next_states = []
self.dones = []
def reset_noise(self):
if not self.discrete_action:
self.exploration.reset()
def scale_noise(self, scale):
if self.discrete_action:
self.exploration = scale
else:
self.exploration.scale = scale
def act(self, obs, explore=False):
"""
Take a step forward in environment for a minibatch of observations
Inputs:
obs : Observations for this agent
explore (boolean): Whether or not to add exploration noise
Outputs:
action (PyTorch Variable): Actions for this agent
"""
#obs = obs.reshape(1,48)
state = Variable(torch.Tensor(obs),requires_grad=False)
self.policy.eval()
with torch.no_grad():
action = self.policy(state)
self.policy.train()
# continuous action
if explore:
action += Variable(Tensor(self.eps * self.exploration.sample()),requires_grad=False)
action = torch.clamp(action, min=-1, max=1)
return action
def step(self, agent_id, state, action, reward, next_state, done,t_step):
self.states.append(state)
self.actions.append(action)
self.rewards.append(reward)
self.next_states.append(next_state)
self.dones.append(done)
#self.replay_buffer.add(state, action, reward, next_state, done)
if t_step % self.num_history == 0:
# Save experience / reward
self.replay_buffer.add(self.states, self.actions, self.rewards, self.next_states, self.dones)
self.states = []
self.actions = []
self.rewards = []
self.next_states = []
self.dones = []
# Learn, if enough samples are available in memory
if len(self.replay_buffer) > self.batch_size:
obs, acs, rews, next_obs, don = self.replay_buffer.sample()
self.update(agent_id ,obs, acs, rews, next_obs, don,t_step)
def update(self, agent_id, obs, acs, rews, next_obs, dones ,t_step, logger=None):
obs = torch.from_numpy(obs).float()
acs = torch.from_numpy(acs).float()
rews = torch.from_numpy(rews[:,agent_id]).float()
next_obs = torch.from_numpy(next_obs).float()
dones = torch.from_numpy(dones[:,agent_id]).float()
acs = acs.view(-1,2)
# --------- update critic ------------ #
self.critic_optimizer.zero_grad()
all_trgt_acs = self.target_policy(next_obs)
target_value = (rews + self.gamma *
self.target_critic(next_obs,all_trgt_acs) *
(1 - dones))
actual_value = self.critic(obs,acs)
vf_loss = MSELoss(actual_value, target_value.detach())
# Minimize the loss
vf_loss.backward()
#torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 1)
self.critic_optimizer.step()
# --------- update actor --------------- #
self.policy_optimizer.zero_grad()
if self.discrete_action:
curr_pol_out = self.policy(obs)
curr_pol_vf_in = gumbel_softmax(curr_pol_out, hard=True)
else:
curr_pol_out = self.policy(obs)
curr_pol_vf_in = curr_pol_out
pol_loss = -self.critic(obs,curr_pol_vf_in).mean()
#pol_loss += (curr_pol_out**2).mean() * 1e-3
pol_loss.backward()
torch.nn.utils.clip_grad_norm_(self.policy.parameters(), 1)
self.policy_optimizer.step()
self.update_all_targets()
self.eps -= self.eps_decay
self.eps = max(self.eps, 0)
if logger is not None:
logger.add_scalars('agent%i/losses' % self.agent_name,
{'vf_loss': vf_loss,
'pol_loss': pol_loss},
self.niter)
def update_all_targets(self):
"""
Update all target networks (called after normal updates have been
performed for each agent)
"""
soft_update(self.critic, self.target_critic, self.tau)
soft_update(self.policy, self.target_policy, self.tau)
def get_params(self):
return {'policy': self.policy.state_dict(),
'critic': self.critic.state_dict(),
'target_policy': self.target_policy.state_dict(),
'target_critic': self.target_critic.state_dict(),
'policy_optimizer': self.policy_optimizer.state_dict(),
'critic_optimizer': self.critic_optimizer.state_dict()}
def load_params(self, params):
self.policy.load_state_dict(params['policy'])
self.critic.load_state_dict(params['critic'])
self.target_policy.load_state_dict(params['target_policy'])
self.target_critic.load_state_dict(params['target_critic'])
self.policy_optimizer.load_state_dict(params['policy_optimizer'])
self.critic_optimizer.load_state_dict(params['critic_optimizer'])
class DDPG(object):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
veer = 'straight'
origin = 'west'
# state_dim : 状态维度
# action_dim： 动作维度
# max_action：动作限制向量
def __init__(self, state_dim, action_dim, max_action, origin_str, veer_str,
logger):
# 存在于 GPU 的神经网络
self.actor = Actor(state_dim, action_dim,
max_action).to(self.device) # origin_network
self.actor_target = Actor(state_dim, action_dim,
max_action).to(self.device) # target_network
self.actor_target.load_state_dict(self.actor.state_dict(
)) # initiate actor_target with actor's parameters
# pytorch 中的 tensor 默认requires_grad 属性为false，即不参与梯度传播运算，特别地，opimizer中模型参数是会参与梯度优化的
self.actor_optimizer = optim.Adam(
self.actor.parameters(),
pdata.LEARNING_RATE) # 以pdata.LEARNING_RATE指定学习率优化actor中的参数
self.critic = Critic(state_dim, action_dim).to(self.device)
self.critic_target = Critic(state_dim, action_dim).to(self.device)
self.critic_target.load_state_dict(self.critic.state_dict())
self.critic_optimizer = optim.Adam(self.critic.parameters(),
pdata.LEARNING_RATE)
# self.replay_buffer = Replay_buffer() # initiate replay-buffer
self.replay_buffer = FilterReplayBuffer()
self.writer = SummaryWriter(pdata.DIRECTORY + 'runs')
self.num_critic_update_iteration = 0
self.num_actor_update_iteration = 0
self.num_training = 0
self.veer = veer_str
self.origin = origin_str
self.logger = logger
def select_action(self, state):
state = torch.FloatTensor(state.reshape(1, -1)).to(self.device)
# numpy.ndarray.flatten(): 返回一个 ndarray对象的copy，并且将该ndarray压缩成一维数组
action = self.actor(state).cpu().data.numpy().flatten()
return action
# update the parameters in actor network and critic network
# 只有 replay_buffer 中的 storage 超过了样本数量才会调用 update函数
def update(self):
critic_loss_list = []
actor_performance_list = []
for it in range(pdata.UPDATE_ITERATION):
# Sample replay buffer
x, y, u, r, d = self.replay_buffer.sample(
pdata.BATCH_SIZE
) # 随机获取 batch_size 个五元组样本(sample random minibatch)
state = torch.FloatTensor(x).to(self.device)
action = torch.FloatTensor(u).to(self.device)
next_state = torch.FloatTensor(y).to(self.device)
done = torch.FloatTensor(d).to(self.device)
reward = torch.FloatTensor(r).to(self.device)
# Compute the target Q value —— Q(S', A') is an value evaluated with next_state and predicted action
# 这里的 target_Q 是 sample 个 一维tensor
target_Q = self.critic_target(next_state,
self.actor_target(next_state))
# detach(): Return a new tensor, detached from the current graph
# evaluate Q: targetQ = R + γQ'(s', a')
target_Q = reward + (
(1 - done) * pdata.GAMMA * target_Q).detach() # batch_size个维度
# Get current Q estimate
current_Q = self.critic(state, action) # 1 维
# Compute critic loss : a mean-square error
# 由论文，critic_loss 其实计算的是每个样本估计值与每个critic网络输出的均值方差
# torch.nn.functional.mse_loss 为计算tensor中各个元素的的均值方差
critic_loss = F.mse_loss(current_Q, target_Q)
self.writer.add_scalar(
'critic_loss',
critic_loss,
global_step=self.num_critic_update_iteration)
# self.logger.write_to_log('critic_loss:{loss}'.format(loss=critic_loss))
# self.logger.add_to_critic_buffer(critic_loss.item())
critic_loss_list.append(critic_loss.item())
# Optimize the critic
self.critic_optimizer.zero_grad() # zeros the gradient buffer
critic_loss.backward() # back propagation on a dynamic graph
self.critic_optimizer.step()
# Compute actor loss
# actor_loss：见论文中对公式 (6) 的理解
# mean()：对tensor对象求所有element的均值
# backward() 以梯度下降的方式更新参数，则将 actor_loss 设置为反向梯度，这样参数便往梯度上升方向更新
actor_loss = -self.critic(state, self.actor(state)).mean()
self.writer.add_scalar('actor_performance',
actor_loss,
global_step=self.num_actor_update_iteration)
# self.logger.write_to_log('actor_loss:{loss}'.format(loss=actor_loss))
# self.logger.add_to_actor_buffer(actor_loss.item())
actor_performance_list.append(actor_loss.item())
# Optimize the actor
self.actor_optimizer.zero_grad(
) # Clears the gradients of all optimized torch.Tensor
actor_loss.backward()
self.actor_optimizer.step() # perform a single optimization step
# 这里是两个 target网络的 soft update
for param, target_param in zip(self.critic.parameters(),
self.critic_target.parameters()):
target_param.data.copy_(pdata.TAU * param.data +
(1 - pdata.TAU) * target_param.data)
for param, target_param in zip(self.actor.parameters(),
self.actor_target.parameters()):
target_param.data.copy_(pdata.TAU * param.data +
(1 - pdata.TAU) * target_param.data)
self.num_actor_update_iteration += 1
self.num_critic_update_iteration += 1
actor_performance = np.mean(np.array(actor_performance_list)).item()
self.logger.add_to_actor_buffer(actor_performance)
critic_loss = np.mean(critic_loss_list).item()
self.logger.add_to_critic_buffer(critic_loss)
def save(self, mark_str):
torch.save(self.actor.state_dict(),
pdata.DIRECTORY + mark_str + '_actor.pth')
torch.save(self.critic.state_dict(),
pdata.DIRECTORY + mark_str + '_critic.pth')
def load(self, mark_str):
file_actor = pdata.DIRECTORY + mark_str + '_actor.pth'
file_critic = pdata.DIRECTORY + mark_str + '_critic.pth'
if os.path.exists(file_actor) and os.path.exists(file_critic):
self.actor.load_state_dict(torch.load(file_actor))
self.critic.load_state_dict(torch.load(file_critic))
class DDPGAgent:
def __init__(self, state_size=24, action_size=2, seed=1, num_agents=2):
self.state_size = state_size
self.action_size = action_size
self.seed = random.seed(seed)
self.num_agents = num_agents
# DDPG specific configuration
hidden_size = 512
self.CHECKPOINT_FOLDER = './'
# Defining networks
self.actor = Actor(state_size, hidden_size, action_size).to(device)
self.actor_target = Actor(state_size, hidden_size, action_size).to(device)
self.critic = Critic(state_size, self.action_size, hidden_size, 1).to(device)
self.critic_target = Critic(state_size, self.action_size, hidden_size, 1).to(device)
self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=ACTOR_LR)
self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=CRITIC_LR)
# Noise
self.noises = OUNoise((num_agents, action_size), seed)
# Initialize replay buffer
self.memory = ReplayBuffer(action_size, BUFFER_SIZE, BATCH_SIZE, seed)
def act(self, state, add_noise=True):
'''
Returns action to be taken based on state provided as the input
'''
state = torch.from_numpy(state).float().to(device)
actions = np.zeros((self.num_agents, self.action_size))
self.actor.eval()
with torch.no_grad():
for agent_num, state in enumerate(state):
action = self.actor(state).cpu().data.numpy()
actions[agent_num, :] = action
self.actor.train()
if add_noise:
actions += self.noises.sample()
return np.clip(actions, -1, 1)
def reset(self):
self.noises.reset()
def learn(self, experiences):
'''
Trains the actor critic network using experiences
'''
states, actions, rewards, next_states, dones = experiences
# Update Critic
actions_next = self.actor_target(next_states)
# print(next_states.shape, actions_next.shape)
Q_targets_next = self.critic_target(next_states, actions_next)
Q_targets = rewards + GAMMA*Q_targets_next*(1-dones)
Q_expected = self.critic(states, actions)
critic_loss = F.mse_loss(Q_expected, Q_targets)
self.optimizer_critic.zero_grad()
critic_loss.backward()
self.optimizer_critic.step()
# Update Actor
actions_pred = self.actor(states)
actor_loss = -self.critic(states, actions_pred).mean()
self.optimizer_actor.zero_grad()
actor_loss.backward()
self.optimizer_actor.step()
# Updating the local networks
self.soft_update(self.critic, self.critic_target)
self.soft_update(self.actor, self.actor_target)
def soft_update(self, model, model_target):
tau = TAU
for target_param, local_param in zip(model_target.parameters(), model.parameters()):
target_param.data.copy_(tau * local_param.data + (1.0 - tau) * target_param.data)
def step(self, state, action, reward, next_state, done):
'''
Adds experience to memory and learns if the memory contains sufficient samples
'''
for i in range(self.num_agents):
self.memory.add(state[i, :], action[i, :], reward[i], next_state[i, :], done[i])
if len(self.memory) > BATCH_SIZE:
# print("Now Learning")
experiences = self.memory.sample()
self.learn(experiences)
def checkpoint(self):
'''
Saves the actor critic network on disk
'''
torch.save(self.actor.state_dict(), self.CHECKPOINT_FOLDER + 'checkpoint_actor.pth')
torch.save(self.critic.state_dict(), self.CHECKPOINT_FOLDER + 'checkpoint_critic.pth')
def load(self):
'''
Loads the actor critic network from disk
'''
self.actor.load_state_dict(torch.load(self.CHECKPOINT_FOLDER + 'checkpoint_actor.pth'))
self.actor_target.load_state_dict(torch.load(self.CHECKPOINT_FOLDER + 'checkpoint_actor.pth'))
self.critic.load_state_dict(torch.load(self.CHECKPOINT_FOLDER + 'checkpoint_critic.pth'))
self.critic_target.load_state_dict(torch.load(self.CHECKPOINT_FOLDER + 'checkpoint_critic.pth'))
class SAC:
def __init__(self,
env,
lr=3e-4,
gamma=0.99,
polyak=5e-3,
alpha=0.2,
reward_scale=1.0,
cuda=True,
writer=None):
state_size = env.observation_space.shape[0]
action_size = env.action_space.shape[0]
self.actor = Actor(state_size, action_size)
self.critic = Critic(state_size, action_size)
self.target_critic = Critic(state_size, action_size).eval()
self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr)
self.q1_optimizer = optim.Adam(self.critic.q1.parameters(), lr=lr)
self.q2_optimizer = optim.Adam(self.critic.q2.parameters(), lr=lr)
self.target_critic.load_state_dict(self.critic.state_dict())
for param in self.target_critic.parameters():
param.requires_grad = False
self.memory = ReplayMemory()
self.gamma = gamma
self.alpha = alpha
self.polyak = polyak # Always between 0 and 1, usually close to 1
self.reward_scale = reward_scale
self.writer = writer
self.cuda = cuda
if cuda:
self.actor = self.actor.to('cuda')
self.critic = self.critic.to('cuda')
self.target_critic = self.target_critic.to('cuda')
def explore(self, state):
if self.cuda:
state = torch.tensor(state).unsqueeze(0).to('cuda', torch.float)
action, _, _ = self.actor.sample(state)
# action, _ = self.actor(state)
return action.cpu().detach().numpy().reshape(-1)
def exploit(self, state):
if self.cuda:
state = torch.tensor(state).unsqueeze(0).to('cuda', torch.float)
_, _, action = self.actor.sample(state)
return action.cpu().detach().numpy().reshape(-1)
def store_step(self, state, action, next_state, reward, terminal):
state = to_tensor_unsqueeze(state)
if action.dtype == np.float32:
action = torch.from_numpy(action)
next_state = to_tensor_unsqueeze(next_state)
reward = torch.from_numpy(np.array([reward]).astype(np.float))
terminal = torch.from_numpy(np.array([terminal]).astype(np.uint8))
self.memory.push(state, action, next_state, reward, terminal)
def target_update(self, target_net, net):
for t, s in zip(target_net.parameters(), net.parameters()):
# t.data.copy_(t.data * (1.0 - self.polyak) + s.data * self.polyak)
t.data.mul_(1.0 - self.polyak)
t.data.add_(self.polyak * s.data)
def calc_target_q(self, next_states, rewards, terminals):
with torch.no_grad():
next_action, entropy, _ = self.actor.sample(
next_states) # penalty term
next_q1, next_q2 = self.target_critic(next_states, next_action)
next_q = torch.min(next_q1, next_q2) - self.alpha * entropy
target_q = rewards * self.reward_scale + (
1. - terminals) * self.gamma * next_q
return target_q
def calc_critic_loss(self, states, actions, next_states, rewards,
terminals):
q1, q2 = self.critic(states, actions)
target_q = self.calc_target_q(next_states, rewards, terminals)
q1_loss = torch.mean((q1 - target_q).pow(2))
q2_loss = torch.mean((q2 - target_q).pow(2))
return q1_loss, q2_loss
def calc_actor_loss(self, states):
action, entropy, _ = self.actor.sample(states)
q1, q2 = self.critic(states, action)
q = torch.min(q1, q2)
# actor_loss = torch.mean(-q - self.alpha * entropy)
actor_loss = (self.alpha * entropy - q).mean()
return actor_loss, entropy
def train(self, timestep, batch_size=256):
if len(self.memory) < batch_size:
return
transitions = self.memory.sample(batch_size)
transitions = Transition(*zip(*transitions))
if self.cuda:
states = torch.cat(transitions.state).to('cuda')
actions = torch.stack(transitions.action).to('cuda')
next_states = torch.cat(transitions.next_state).to('cuda')
rewards = torch.stack(transitions.reward).to('cuda')
terminals = torch.stack(transitions.terminal).to('cuda')
else:
states = torch.cat(transitions.state)
actions = torch.stack(transitions.action)
next_states = torch.cat(transitions.next_state)
rewards = torch.stack(transitions.reward)
terminals = torch.stack(transitions.terminal)
# Compute target Q func
q1_loss, q2_loss = self.calc_critic_loss(states, actions, next_states,
rewards, terminals)
# Compute actor loss
actor_loss, mean = self.calc_actor_loss(states)
update_params(self.q1_optimizer, self.critic.q1, q1_loss)
update_params(self.q2_optimizer, self.critic.q2, q2_loss)
update_params(self.actor_optimizer, self.actor, actor_loss)
# target update
self.target_update(self.target_critic, self.critic)
if timestep % 100 and self.writer:
self.writer.add_scalar('Loss/Actor', actor_loss.item(), timestep)
self.writer.add_scalar('Loss/Critic', q1_loss.item(), timestep)
def save_weights(self, path):
self.actor.save(os.path.join(path, 'actor'))
self.critic.save(os.path.join(path, 'critic'))