mirror of
https://github.com/wassname/kair_algorithms_draft.git
synced 2026-09-09 11:25:10 +08:00
Add DDPGfD, TD3fD and SACfD (#22)
* Format repository * Clone files from medipixel repo * Fix DDPGfDAgent.update_model() * Fix bug on _initialize() * Add demo-path parameter and demo data * Rename init_priority to _max_priority for PER This makes PER and PERfD consistent. * Make i_episode attribute of DDPGAgent * Clone SAC code from medipixel repo * Fix update_model() for SACfD * Fix _initialize() for SACfD * Add is_discrete attribute to AbstractAgent for SACfD * Add i_episode attribute to SACAgent for SACfD * Modularize DDPGAgent and SACAgent * Modify hyperparameters for DDPGfD and SACfD * Add NStepBuffer * Add n-step to DDPGfD * Add n-step to SACfD * Add TD3fD without n-step * Attempt to tune hyperparameters * Remove discrete environment check in SAC * Implement n-step on TD3fD * Fix step function of TD3 No done check, and _add_transition_to_memory was not called. * Fix actor loss calculation for TD3fD * Attempt to tune hyperparameters * Print both critic losses * Fix typo bug * Attempt to tune hyperparameters * Fix bug in n-step demo retrieval * Fix bug in n-step transition addition
This commit is contained in:
Seungjae Ryan Lee
@@ -0,0 +1,176 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""DDPGfD agent using demo agent for episodic tasks in OpenAI Gym.
|
||||
|
||||
- Author: Kh Kim
|
||||
- Contact: kh.kim@medipixel.io
|
||||
- Paper: https://arxiv.org/pdf/1509.02971.pdf
|
||||
https://arxiv.org/pdf/1511.05952.pdf
|
||||
https://arxiv.org/pdf/1707.08817.pdf
|
||||
"""
|
||||
|
||||
import pickle
|
||||
from typing import List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import algorithms.common.helper_functions as common_utils
|
||||
from algorithms.common.buffer.priortized_replay_buffer import PrioritizedReplayBufferfD
|
||||
from algorithms.common.buffer.replay_buffer import NStepTransitionBuffer
|
||||
from algorithms.ddpg.agent import Agent as DDPGAgent
|
||||
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class Agent(DDPGAgent):
|
||||
"""ActorCritic interacting with environment.
|
||||
|
||||
Attributes:
|
||||
memory (PrioritizedReplayBufferfD): replay memory
|
||||
beta (float): beta parameter for prioritized replay buffer
|
||||
|
||||
"""
|
||||
|
||||
# pylint: disable=attribute-defined-outside-init
|
||||
def _initialize(self):
|
||||
"""Initialize non-common things."""
|
||||
self.use_n_step = self.hyper_params["N_STEP"] > 1
|
||||
|
||||
if not self.args.test:
|
||||
# load demo replay memory
|
||||
with open(self.args.demo_path, "rb") as f:
|
||||
demos = pickle.load(f)
|
||||
|
||||
if self.use_n_step:
|
||||
demos, demos_n_step = common_utils.get_n_step_info_from_demo(
|
||||
demos, self.hyper_params["N_STEP"], self.hyper_params["GAMMA"]
|
||||
)
|
||||
|
||||
# replay memory for multi-steps
|
||||
self.memory_n = NStepTransitionBuffer(
|
||||
buffer_size=self.hyper_params["BUFFER_SIZE"],
|
||||
n_step=self.hyper_params["N_STEP"],
|
||||
gamma=self.hyper_params["GAMMA"],
|
||||
demo=demos_n_step,
|
||||
)
|
||||
|
||||
# replay memory for a single step
|
||||
self.beta = self.hyper_params["PER_BETA"]
|
||||
self.memory = PrioritizedReplayBufferfD(
|
||||
self.hyper_params["BUFFER_SIZE"],
|
||||
self.hyper_params["BATCH_SIZE"],
|
||||
demo=list(demos),
|
||||
alpha=self.hyper_params["PER_ALPHA"],
|
||||
epsilon_d=self.hyper_params["PER_EPS_DEMO"],
|
||||
)
|
||||
|
||||
def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]):
|
||||
"""Add 1 step and n step transitions to memory."""
|
||||
# add n-step transition
|
||||
if self.use_n_step:
|
||||
transition = self.memory_n.add(transition)
|
||||
|
||||
# add a single step transition
|
||||
# if transition is not an empty tuple
|
||||
if transition:
|
||||
self.memory.add(*transition)
|
||||
|
||||
def _get_critic_loss(
|
||||
self, experiences: Tuple[torch.Tensor, ...], gamma: float
|
||||
) -> torch.Tensor:
|
||||
"""Return element-wise critic loss."""
|
||||
states, actions, rewards, next_states, dones = experiences[:5]
|
||||
|
||||
# G_t = r + gamma * v(s_{t+1}) if state != Terminal
|
||||
# = r otherwise
|
||||
masks = 1 - dones
|
||||
next_actions = self.actor_target(next_states)
|
||||
next_states_actions = torch.cat((next_states, next_actions), dim=-1)
|
||||
next_values = self.critic_target(next_states_actions)
|
||||
curr_returns = rewards + gamma * next_values * masks
|
||||
curr_returns = curr_returns.to(device).detach()
|
||||
|
||||
# train critic
|
||||
values = self.critic(torch.cat((states, actions), dim=-1))
|
||||
critic_loss_element_wise = (values - curr_returns).pow(2)
|
||||
|
||||
return critic_loss_element_wise
|
||||
|
||||
def update_model(
|
||||
self,
|
||||
experiences: Tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
List[int],
|
||||
],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Train the model after each episode."""
|
||||
# NOTE This is for old update_model() interface.
|
||||
# experiences_1 = self.memory.sample(self.beta)
|
||||
experiences_1 = experiences
|
||||
states, actions = experiences_1[:2]
|
||||
weights, indices, eps_d = experiences_1[-3:]
|
||||
gamma = self.hyper_params["GAMMA"]
|
||||
|
||||
# train critic
|
||||
critic_loss_element_wise = self._get_critic_loss(experiences_1, gamma)
|
||||
critic_loss = torch.mean(critic_loss_element_wise * weights)
|
||||
|
||||
if self.use_n_step:
|
||||
experiences_n = self.memory_n.sample(indices)
|
||||
gamma = gamma ** self.hyper_params["N_STEP"]
|
||||
critic_loss_n_element_wise = self._get_critic_loss(experiences_n, gamma)
|
||||
# to update loss and priorities
|
||||
lambda1 = self.hyper_params["LAMBDA1"]
|
||||
critic_loss_element_wise += critic_loss_n_element_wise * lambda1
|
||||
critic_loss = torch.mean(critic_loss_element_wise * weights)
|
||||
|
||||
self.critic_optimizer.zero_grad()
|
||||
critic_loss.backward()
|
||||
self.critic_optimizer.step()
|
||||
|
||||
# train actor
|
||||
actions = self.actor(states)
|
||||
actor_loss_element_wise = -self.critic(torch.cat((states, actions), dim=-1))
|
||||
actor_loss = torch.mean(actor_loss_element_wise * weights)
|
||||
self.actor_optimizer.zero_grad()
|
||||
actor_loss.backward()
|
||||
self.actor_optimizer.step()
|
||||
|
||||
# update target networks
|
||||
tau = self.hyper_params["TAU"]
|
||||
common_utils.soft_update(self.actor, self.actor_target, tau)
|
||||
common_utils.soft_update(self.critic, self.critic_target, tau)
|
||||
|
||||
# update priorities
|
||||
new_priorities = critic_loss_element_wise
|
||||
new_priorities += self.hyper_params["LAMBDA3"] * actor_loss_element_wise.pow(2)
|
||||
new_priorities += self.hyper_params["PER_EPS"]
|
||||
new_priorities = new_priorities.data.cpu().numpy().squeeze()
|
||||
new_priorities += eps_d
|
||||
self.memory.update_priorities(indices, new_priorities)
|
||||
|
||||
# increase beta
|
||||
fraction = min(float(self.i_episode) / self.args.episode_num, 1.0)
|
||||
self.beta = self.beta + fraction * (1.0 - self.beta)
|
||||
|
||||
return actor_loss.data, critic_loss.data
|
||||
|
||||
def pretrain(self):
|
||||
"""Pretraining steps."""
|
||||
pretrain_loss = list()
|
||||
print("[INFO] Pre-Train %d step." % self.hyper_params["PRETRAIN_STEP"])
|
||||
for i_step in range(1, self.hyper_params["PRETRAIN_STEP"] + 1):
|
||||
loss = self.update_model()
|
||||
pretrain_loss.append(loss) # for logging
|
||||
|
||||
# logging
|
||||
if i_step == 1 or i_step % 100 == 0:
|
||||
avg_loss = np.vstack(pretrain_loss).mean(axis=0)
|
||||
pretrain_loss.clear()
|
||||
self.write_log(0, avg_loss, 0)
|
||||
@@ -0,0 +1,226 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""SAC agent from demonstration for episodic tasks in OpenAI Gym.
|
||||
|
||||
- Author: Curt Park
|
||||
- Contact: curt.park@medipixel.io
|
||||
- Paper: https://arxiv.org/pdf/1801.01290.pdf
|
||||
https://arxiv.org/pdf/1812.05905.pdf
|
||||
https://arxiv.org/pdf/1511.05952.pdf
|
||||
https://arxiv.org/pdf/1707.08817.pdf
|
||||
"""
|
||||
|
||||
import pickle
|
||||
from typing import List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import algorithms.common.helper_functions as common_utils
|
||||
from algorithms.common.buffer.priortized_replay_buffer import PrioritizedReplayBufferfD
|
||||
from algorithms.common.buffer.replay_buffer import NStepTransitionBuffer
|
||||
from algorithms.sac.agent import Agent as SACAgent
|
||||
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class Agent(SACAgent):
|
||||
"""SAC agent interacting with environment.
|
||||
|
||||
Attrtibutes:
|
||||
memory (PrioritizedReplayBufferfD): replay memory
|
||||
beta (float): beta parameter for prioritized replay buffer
|
||||
|
||||
"""
|
||||
|
||||
# pylint: disable=attribute-defined-outside-init
|
||||
def _initialize(self):
|
||||
"""Initialize non-common things."""
|
||||
self.use_n_step = self.hyper_params["N_STEP"] > 1
|
||||
|
||||
if not self.args.test:
|
||||
# load demo replay memory
|
||||
with open(self.args.demo_path, "rb") as f:
|
||||
demos = pickle.load(f)
|
||||
|
||||
if self.use_n_step:
|
||||
demos, demos_n_step = common_utils.get_n_step_info_from_demo(
|
||||
demos, self.hyper_params["N_STEP"], self.hyper_params["GAMMA"]
|
||||
)
|
||||
|
||||
# replay memory for multi-steps
|
||||
self.memory_n = NStepTransitionBuffer(
|
||||
buffer_size=self.hyper_params["BUFFER_SIZE"],
|
||||
n_step=self.hyper_params["N_STEP"],
|
||||
gamma=self.hyper_params["GAMMA"],
|
||||
demo=demos_n_step,
|
||||
)
|
||||
|
||||
# replay memory
|
||||
self.beta = self.hyper_params["PER_BETA"]
|
||||
self.memory = PrioritizedReplayBufferfD(
|
||||
self.hyper_params["BUFFER_SIZE"],
|
||||
self.hyper_params["BATCH_SIZE"],
|
||||
demo=demos,
|
||||
alpha=self.hyper_params["PER_ALPHA"],
|
||||
epsilon_d=self.hyper_params["PER_EPS_DEMO"],
|
||||
)
|
||||
|
||||
def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]):
|
||||
"""Add 1 step and n step transitions to memory."""
|
||||
# add n-step transition
|
||||
if self.use_n_step:
|
||||
transition = self.memory_n.add(transition)
|
||||
|
||||
# add a single step transition
|
||||
# if transition is not an empty tuple
|
||||
if transition:
|
||||
self.memory.add(*transition)
|
||||
|
||||
# pylint: disable=too-many-statements
|
||||
def update_model(
|
||||
self,
|
||||
experiences: Tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
List[int],
|
||||
],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Train the model after each episode."""
|
||||
states, actions, rewards, next_states, dones, weights, indices, eps_d = (
|
||||
experiences
|
||||
)
|
||||
new_actions, log_prob, pre_tanh_value, mu, std = self.actor(states)
|
||||
|
||||
# train alpha
|
||||
if self.hyper_params["AUTO_ENTROPY_TUNING"]:
|
||||
alpha_loss = torch.mean(
|
||||
(-self.log_alpha * (log_prob + self.target_entropy).detach()) * weights
|
||||
)
|
||||
|
||||
self.alpha_optimizer.zero_grad()
|
||||
alpha_loss.backward()
|
||||
self.alpha_optimizer.step()
|
||||
|
||||
alpha = self.log_alpha.exp()
|
||||
else:
|
||||
alpha_loss = torch.zeros(1)
|
||||
alpha = self.hyper_params["W_ENTROPY"]
|
||||
|
||||
# Q function loss
|
||||
masks = 1 - dones
|
||||
gamma = self.hyper_params["GAMMA"]
|
||||
q_1_pred = self.qf_1(states, actions)
|
||||
q_2_pred = self.qf_2(states, actions)
|
||||
v_target = self.vf_target(next_states)
|
||||
q_target = rewards + self.hyper_params["GAMMA"] * v_target * masks
|
||||
qf_1_loss = torch.mean((q_1_pred - q_target.detach()).pow(2) * weights)
|
||||
qf_2_loss = torch.mean((q_2_pred - q_target.detach()).pow(2) * weights)
|
||||
|
||||
if self.use_n_step:
|
||||
experiences_n = self.memory_n.sample(indices)
|
||||
_, _, rewards, next_states, dones = experiences_n
|
||||
gamma = gamma ** self.hyper_params["N_STEP"]
|
||||
lambda1 = self.hyper_params["LAMBDA1"]
|
||||
masks = 1 - dones
|
||||
|
||||
v_target = self.vf_target(next_states)
|
||||
q_target = rewards + gamma * v_target * masks
|
||||
qf_1_loss_n = torch.mean((q_1_pred - q_target.detach()).pow(2) * weights)
|
||||
qf_2_loss_n = torch.mean((q_2_pred - q_target.detach()).pow(2) * weights)
|
||||
|
||||
# to update loss and priorities
|
||||
qf_1_loss = qf_1_loss + qf_1_loss_n * lambda1
|
||||
qf_2_loss = qf_2_loss + qf_2_loss_n * lambda1
|
||||
|
||||
# V function loss
|
||||
v_pred = self.vf(states)
|
||||
q_pred = torch.min(
|
||||
self.qf_1(states, new_actions), self.qf_2(states, new_actions)
|
||||
)
|
||||
v_target = (q_pred - alpha * log_prob).detach()
|
||||
vf_loss_element_wise = (v_pred - v_target).pow(2)
|
||||
vf_loss = torch.mean(vf_loss_element_wise * weights)
|
||||
|
||||
# train Q functions
|
||||
self.qf_1_optimizer.zero_grad()
|
||||
qf_1_loss.backward()
|
||||
self.qf_1_optimizer.step()
|
||||
|
||||
self.qf_2_optimizer.zero_grad()
|
||||
qf_2_loss.backward()
|
||||
self.qf_2_optimizer.step()
|
||||
|
||||
# train V function
|
||||
self.vf_optimizer.zero_grad()
|
||||
vf_loss.backward()
|
||||
self.vf_optimizer.step()
|
||||
|
||||
if self.total_step % self.hyper_params["DELAYED_UPDATE"] == 0:
|
||||
# actor loss
|
||||
advantage = q_pred - v_pred.detach()
|
||||
actor_loss_element_wise = alpha * log_prob - advantage
|
||||
actor_loss = torch.mean(actor_loss_element_wise * weights)
|
||||
|
||||
# regularization
|
||||
mean_reg = self.hyper_params["W_MEAN_REG"] * mu.pow(2).mean()
|
||||
std_reg = self.hyper_params["W_STD_REG"] * std.pow(2).mean()
|
||||
pre_activation_reg = self.hyper_params["W_PRE_ACTIVATION_REG"] * (
|
||||
pre_tanh_value.pow(2).sum(dim=-1).mean()
|
||||
)
|
||||
actor_reg = mean_reg + std_reg + pre_activation_reg
|
||||
|
||||
# actor loss + regularization
|
||||
actor_loss += actor_reg
|
||||
|
||||
# train actor
|
||||
self.actor_optimizer.zero_grad()
|
||||
actor_loss.backward()
|
||||
self.actor_optimizer.step()
|
||||
|
||||
# update target networks
|
||||
common_utils.soft_update(self.vf, self.vf_target, self.hyper_params["TAU"])
|
||||
|
||||
# update priorities
|
||||
new_priorities = vf_loss_element_wise
|
||||
new_priorities += self.hyper_params[
|
||||
"LAMBDA3"
|
||||
] * actor_loss_element_wise.pow(2)
|
||||
new_priorities += self.hyper_params["PER_EPS"]
|
||||
new_priorities = new_priorities.data.cpu().numpy().squeeze()
|
||||
new_priorities += eps_d
|
||||
self.memory.update_priorities(indices, new_priorities)
|
||||
|
||||
# increase beta
|
||||
fraction = min(float(self.i_episode) / self.args.episode_num, 1.0)
|
||||
self.beta = self.beta + fraction * (1.0 - self.beta)
|
||||
else:
|
||||
actor_loss = torch.zeros(1)
|
||||
|
||||
return (
|
||||
actor_loss.data,
|
||||
qf_1_loss.data,
|
||||
qf_2_loss.data,
|
||||
vf_loss.data,
|
||||
alpha_loss.data,
|
||||
)
|
||||
|
||||
def pretrain(self):
|
||||
"""Pretraining steps."""
|
||||
pretrain_loss = list()
|
||||
print("[INFO] Pre-Train %d steps." % self.hyper_params["PRETRAIN_STEP"])
|
||||
for i_step in range(1, self.hyper_params["PRETRAIN_STEP"] + 1):
|
||||
loss = self.update_model()
|
||||
pretrain_loss.append(loss) # for logging
|
||||
|
||||
# logging
|
||||
if i_step == 1 or i_step % 100 == 0:
|
||||
avg_loss = np.vstack(pretrain_loss).mean(axis=0)
|
||||
pretrain_loss.clear()
|
||||
self.write_log(
|
||||
0, avg_loss, 0, delayed_update=self.hyper_params["DELAYED_UPDATE"]
|
||||
)
|
||||
@@ -0,0 +1,209 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""TD3 agent from demonstration for episodic tasks in OpenAI Gym.
|
||||
|
||||
- Author: Seungjae Ryan Lee
|
||||
- Contact: seungjaeryanlee@gmail.com
|
||||
- Paper: https://arxiv.org/pdf/1802.09477.pdf (TD3)
|
||||
https://arxiv.org/pdf/1511.05952.pdf (PER)
|
||||
https://arxiv.org/pdf/1707.08817.pdf (DDPGfD)
|
||||
"""
|
||||
|
||||
import pickle
|
||||
from typing import List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import algorithms.common.helper_functions as common_utils
|
||||
from algorithms.common.buffer.priortized_replay_buffer import PrioritizedReplayBufferfD
|
||||
from algorithms.common.buffer.replay_buffer import NStepTransitionBuffer
|
||||
from algorithms.td3.agent import Agent as TD3Agent
|
||||
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class Agent(TD3Agent):
|
||||
"""TD3 agent interacting with environment.
|
||||
|
||||
Attrtibutes:
|
||||
memory (PrioritizedReplayBufferfD): replay memory
|
||||
beta (float): beta parameter for prioritized replay buffer
|
||||
|
||||
"""
|
||||
|
||||
# pylint: disable=attribute-defined-outside-init
|
||||
def _initialize(self):
|
||||
"""Initialize non-common things."""
|
||||
self.use_n_step = self.hyper_params["N_STEP"] > 1
|
||||
|
||||
if not self.args.test:
|
||||
# load demo replay memory
|
||||
with open(self.args.demo_path, "rb") as f:
|
||||
demos = pickle.load(f)
|
||||
|
||||
if self.use_n_step:
|
||||
demos, demos_n_step = common_utils.get_n_step_info_from_demo(
|
||||
demos, self.hyper_params["N_STEP"], self.hyper_params["GAMMA"]
|
||||
)
|
||||
|
||||
# replay memory for multi-steps
|
||||
self.memory_n = NStepTransitionBuffer(
|
||||
buffer_size=self.hyper_params["BUFFER_SIZE"],
|
||||
n_step=self.hyper_params["N_STEP"],
|
||||
gamma=self.hyper_params["GAMMA"],
|
||||
demo=demos_n_step,
|
||||
)
|
||||
|
||||
# replay memory
|
||||
self.beta = self.hyper_params["PER_BETA"]
|
||||
self.memory = PrioritizedReplayBufferfD(
|
||||
self.hyper_params["BUFFER_SIZE"],
|
||||
self.hyper_params["BATCH_SIZE"],
|
||||
demo=demos,
|
||||
alpha=self.hyper_params["PER_ALPHA"],
|
||||
epsilon_d=self.hyper_params["PER_EPS_DEMO"],
|
||||
)
|
||||
|
||||
def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]):
|
||||
"""Add 1 step and n step transitions to memory."""
|
||||
# add n-step transition
|
||||
if self.use_n_step:
|
||||
transition = self.memory_n.add(transition)
|
||||
|
||||
# add a single step transition
|
||||
# if transition is not an empty tuple
|
||||
if transition:
|
||||
self.memory.add(*transition)
|
||||
|
||||
def _get_critic_loss(
|
||||
self, experiences: Tuple[torch.Tensor, ...], gamma: float
|
||||
) -> torch.Tensor:
|
||||
"""Return element-wise critic loss."""
|
||||
states, actions, rewards, next_states, dones = experiences[:5]
|
||||
|
||||
# G_t = r + gamma * v(s_{t+1}) if state != Terminal
|
||||
# = r otherwise
|
||||
masks = 1 - dones
|
||||
noise = torch.FloatTensor(self.target_policy_noise.sample()).to(device)
|
||||
clipped_noise = torch.clamp(
|
||||
noise,
|
||||
-self.hyper_params["TARGET_POLICY_NOISE_CLIP"],
|
||||
self.hyper_params["TARGET_POLICY_NOISE_CLIP"],
|
||||
)
|
||||
next_actions = (self.actor_target(next_states) + clipped_noise).clamp(-1.0, 1.0)
|
||||
|
||||
target_values1 = self.critic1_target(
|
||||
torch.cat((next_states, next_actions), dim=-1)
|
||||
)
|
||||
target_values2 = self.critic2_target(
|
||||
torch.cat((next_states, next_actions), dim=-1)
|
||||
)
|
||||
target_values = torch.min(target_values1, target_values2)
|
||||
target_values = (
|
||||
rewards + (self.hyper_params["GAMMA"] * target_values * masks).detach()
|
||||
)
|
||||
|
||||
# train critic
|
||||
values1 = self.critic1(torch.cat((states, actions), dim=-1))
|
||||
critic1_loss_element_wise = (values1 - target_values.detach()).pow(2)
|
||||
|
||||
values2 = self.critic2(torch.cat((states, actions), dim=-1))
|
||||
critic2_loss_element_wise = (values2 - target_values.detach()).pow(2)
|
||||
|
||||
return critic1_loss_element_wise, critic2_loss_element_wise
|
||||
|
||||
# pylint: disable=too-many-statements
|
||||
def update_model(
|
||||
self,
|
||||
experiences: Tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
List[int],
|
||||
],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Train the model after each episode."""
|
||||
states, actions, rewards, next_states, dones, weights, indices, eps_d = (
|
||||
experiences
|
||||
)
|
||||
|
||||
gamma = self.hyper_params["GAMMA"]
|
||||
critic1_loss_element_wise, critic2_loss_element_wise = self._get_critic_loss(
|
||||
experiences, gamma
|
||||
)
|
||||
critic_loss_element_wise = critic1_loss_element_wise + critic2_loss_element_wise
|
||||
critic1_loss = torch.mean(critic1_loss_element_wise * weights)
|
||||
critic2_loss = torch.mean(critic2_loss_element_wise * weights)
|
||||
critic_loss = critic1_loss + critic2_loss
|
||||
|
||||
if self.use_n_step:
|
||||
experiences_n = self.memory_n.sample(indices)
|
||||
gamma = self.hyper_params["GAMMA"] ** self.hyper_params["N_STEP"]
|
||||
critic1_loss_n_element_wise, critic2_loss_n_element_wise = self._get_critic_loss(
|
||||
experiences_n, gamma
|
||||
)
|
||||
critic_loss_n_element_wise = (
|
||||
critic1_loss_n_element_wise + critic2_loss_n_element_wise
|
||||
)
|
||||
critic1_loss_n = torch.mean(critic1_loss_n_element_wise * weights)
|
||||
critic2_loss_n = torch.mean(critic2_loss_n_element_wise * weights)
|
||||
critic_loss_n = critic1_loss_n + critic2_loss_n
|
||||
|
||||
lambda1 = self.hyper_params["LAMBDA1"]
|
||||
critic_loss_element_wise += lambda1 * critic_loss_n_element_wise
|
||||
critic_loss += lambda1 * critic_loss_n
|
||||
|
||||
self.critic_optim.zero_grad()
|
||||
critic_loss.backward()
|
||||
self.critic_optim.step()
|
||||
|
||||
if self.episode_steps % self.hyper_params["POLICY_UPDATE_FREQ"] == 0:
|
||||
# train actor
|
||||
actions = self.actor(states)
|
||||
actor_loss_element_wise = -self.critic1(
|
||||
torch.cat((states, actions), dim=-1)
|
||||
)
|
||||
actor_loss = torch.mean(actor_loss_element_wise * weights)
|
||||
self.actor_optim.zero_grad()
|
||||
actor_loss.backward()
|
||||
self.actor_optim.step()
|
||||
|
||||
# update target networks
|
||||
tau = self.hyper_params["TAU"]
|
||||
common_utils.soft_update(self.actor, self.actor_target, tau)
|
||||
common_utils.soft_update(self.critic1, self.critic1_target, tau)
|
||||
common_utils.soft_update(self.critic2, self.critic2_target, tau)
|
||||
|
||||
# update priorities
|
||||
new_priorities = critic_loss_element_wise
|
||||
new_priorities += self.hyper_params[
|
||||
"LAMBDA3"
|
||||
] * actor_loss_element_wise.pow(2)
|
||||
new_priorities += self.hyper_params["PER_EPS"]
|
||||
new_priorities = new_priorities.data.cpu().numpy().squeeze()
|
||||
new_priorities += eps_d
|
||||
self.memory.update_priorities(indices, new_priorities)
|
||||
else:
|
||||
actor_loss = torch.zeros(1)
|
||||
|
||||
return actor_loss.data, critic1_loss.data, critic2_loss.data
|
||||
|
||||
def pretrain(self):
|
||||
"""Pretraining steps."""
|
||||
pretrain_loss = list()
|
||||
print("[INFO] Pre-Train %d steps." % self.hyper_params["PRETRAIN_STEP"])
|
||||
for i_step in range(1, self.hyper_params["PRETRAIN_STEP"] + 1):
|
||||
loss = self.update_model()
|
||||
pretrain_loss.append(loss) # for logging
|
||||
|
||||
# logging
|
||||
if i_step == 1 or i_step % 100 == 0:
|
||||
avg_loss = np.vstack(pretrain_loss).mean(axis=0)
|
||||
pretrain_loss.clear()
|
||||
self.write_log(
|
||||
0, avg_loss, 0, delayed_update=self.hyper_params["DELAYED_UPDATE"]
|
||||
)
|
||||
Reference in New Issue
Block a user