mirror of
https://github.com/wassname/kair_algorithms_draft.git
synced 2026-08-21 11:16:38 +08:00
195 lines
7.6 KiB
Python
195 lines
7.6 KiB
Python
# -*- 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
|
|
|
|
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
|
|
# TODO: should make new demo to set protocol 2
|
|
# e.g. pickle.dump(your_object, your_file, protocol=2)
|
|
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):
|
|
"""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, gamma):
|
|
"""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 + (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):
|
|
"""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"]
|
|
)
|