mirror of
https://github.com/wassname/kair_algorithms_draft.git
synced 2026-08-21 11:16:38 +08:00
272 lines
10 KiB
Python
272 lines
10 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""TD3 agent for episodic tasks in OpenAI Gym.
|
|
|
|
- Author: whikwon
|
|
- Contact: whikwon@gmail.com
|
|
- Paper: https://arxiv.org/pdf/1802.09477.pdf
|
|
"""
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import wandb
|
|
|
|
import algorithms.common.helper_functions as common_utils
|
|
from algorithms.common.abstract.agent import AbstractAgent
|
|
from algorithms.common.buffer.replay_buffer import ReplayBuffer
|
|
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
class Agent(AbstractAgent):
|
|
"""ActorCritic interacting with environment.
|
|
|
|
Attributes:
|
|
memory (ReplayBuffer): replay memory
|
|
exploration_noise (GaussianNoise): random noise for exploration
|
|
target_policy_noise (GaussianNoise): random noise for regularization
|
|
hyper_params (dict): hyper-parameters
|
|
actor (nn.Module): actor model to select actions
|
|
actor_target (nn.Module): target actor model to select actions
|
|
critic1 (nn.Module): critic1 model to predict state values
|
|
critic2 (nn.Module): critic2 model to predict state values
|
|
critic1_target (nn.Module): target critic1 model to predict state values
|
|
critic2_target (nn.Module): target critic2 model to predict state values
|
|
actor_optim (Optimizer): optimizer for training actor
|
|
critic1_optim (Optimizer): optimizer for training critic1
|
|
critic2_optim (Optimizer): optimizer for training critic2
|
|
curr_state (np.ndarray): temporary storage of the current state
|
|
|
|
"""
|
|
|
|
def __init__(self, env, args, hyper_params, models, optims, noises):
|
|
"""Initialization.
|
|
|
|
Args:
|
|
env (gym.Env): openAI Gym environment with discrete action space
|
|
args (argparse.Namespace): arguments including hyperparameters and training settings
|
|
hyper_params (dict): hyper-parameters
|
|
models (tuple): models including actor and critics
|
|
optims (tuple): optimizers for actor and critics
|
|
noises (tuple): noises for exploration and regularization
|
|
|
|
"""
|
|
AbstractAgent.__init__(self, env, args)
|
|
self.actor, self.actor_target = models[:2]
|
|
self.critic1, self.critic1_target = models[2:4]
|
|
self.critic2, self.critic2_target = models[4:]
|
|
|
|
self.actor_optim, self.critic_optim = optims
|
|
self.hyper_params = hyper_params
|
|
self.exploration_noise, self.target_policy_noise = noises
|
|
self.curr_state = np.zeros((1,))
|
|
self.total_steps = 0
|
|
self.episode_steps = 0
|
|
|
|
# load the optimizer and model parameters
|
|
if args.load_from is not None and os.path.exists(args.load_from):
|
|
self.load_params(args.load_from)
|
|
|
|
self._initialize()
|
|
|
|
def _initialize(self):
|
|
"""Initialize non-common things."""
|
|
if not self.args.test:
|
|
# replay memory
|
|
self.memory = ReplayBuffer(
|
|
self.hyper_params["BUFFER_SIZE"], self.hyper_params["BATCH_SIZE"]
|
|
)
|
|
|
|
def select_action(self, state):
|
|
"""Select an action from the input space."""
|
|
# initial training step, try random action for exploration
|
|
random_action_count = self.hyper_params["INITIAL_RANDOM_ACTIONS"]
|
|
|
|
self.curr_state = state
|
|
|
|
if self.total_steps < random_action_count and not self.args.test:
|
|
return self.env.action_space.sample()
|
|
|
|
state = torch.FloatTensor(state).to(device)
|
|
selected_action = self.actor(state)
|
|
|
|
if not self.args.test:
|
|
noise = torch.FloatTensor(self.exploration_noise.sample()).to(device)
|
|
selected_action = (selected_action + noise).clamp(-1.0, 1.0)
|
|
|
|
return selected_action.detach().cpu().numpy()
|
|
|
|
def step(self, action):
|
|
"""Take an action and return the response of the env."""
|
|
self.total_steps += 1
|
|
self.episode_steps += 1
|
|
|
|
next_state, reward, done, _ = self.env.step(action)
|
|
|
|
if not self.args.test:
|
|
# if the last state is not a terminal state, store done as false
|
|
done_bool = (
|
|
False if self.episode_steps == self.args.max_episode_steps else done
|
|
)
|
|
transition = (self.curr_state, action, reward, next_state, done_bool)
|
|
self._add_transition_to_memory(transition)
|
|
|
|
return next_state, reward, done
|
|
|
|
def _add_transition_to_memory(self, transition):
|
|
"""Add 1 step and n step transitions to memory."""
|
|
self.memory.add(*transition)
|
|
|
|
def update_model(self, experiences):
|
|
"""Train the model after each episode."""
|
|
states, actions, rewards, next_states, dones = experiences
|
|
|
|
# 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 = F.mse_loss(values1, target_values)
|
|
values2 = self.critic2(torch.cat((states, actions), dim=-1))
|
|
critic2_loss = F.mse_loss(values2, target_values)
|
|
|
|
critic_loss = critic1_loss + critic2_loss
|
|
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 = -self.critic1(torch.cat((states, actions), dim=-1)).mean()
|
|
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)
|
|
else:
|
|
actor_loss = torch.zeros(1)
|
|
|
|
return actor_loss.data, critic1_loss.data, critic2_loss.data
|
|
|
|
def load_params(self, path):
|
|
"""Load model and optimizer parameters."""
|
|
if not os.path.exists(path):
|
|
print("[ERROR] the input path does not exist. ->", path)
|
|
return
|
|
|
|
params = torch.load(path)
|
|
self.actor.load_state_dict(params["actor_state_dict"])
|
|
self.actor_target.load_state_dict(params["actor_target_state_dict"])
|
|
self.critic1.load_state_dict(params["critic1_state_dict"])
|
|
self.critic2.load_state_dict(params["critic2_state_dict"])
|
|
self.critic1_target.load_state_dict(params["critic1_target_state_dict"])
|
|
self.critic2_target.load_state_dict(params["critic2_target_state_dict"])
|
|
self.actor_optim.load_state_dict(params["actor_optim_state_dict"])
|
|
self.critic_optim.load_state_dict(params["critic_optim_state_dict"])
|
|
print("[INFO] loaded the model and optimizer from", path)
|
|
|
|
def save_params(self, n_episode):
|
|
"""Save model and optimizer parameters."""
|
|
params = {
|
|
"actor_state_dict": self.actor.state_dict(),
|
|
"actor_target_state_dict": self.actor_target.state_dict(),
|
|
"critic1_state_dict": self.critic1.state_dict(),
|
|
"critic2_state_dict": self.critic2.state_dict(),
|
|
"critic1_target_state_dict": self.critic1_target.state_dict(),
|
|
"critic2_target_state_dict": self.critic2_target.state_dict(),
|
|
"actor_optim_state_dict": self.actor_optim.state_dict(),
|
|
"critic_optim_state_dict": self.critic_optim.state_dict(),
|
|
}
|
|
|
|
AbstractAgent.save_params(self, params, n_episode)
|
|
|
|
def write_log(self, i, loss, score):
|
|
"""Write log about loss and score"""
|
|
total_loss = loss.sum()
|
|
|
|
print(
|
|
"[INFO] total_steps: %d episode: %d total score: %d, total loss: %f\n"
|
|
"actor_loss: %.3f critic1_loss: %.3f critic2_loss: %.3f\n"
|
|
% (self.total_steps, i, score, total_loss, loss[0], loss[1], loss[2])
|
|
)
|
|
|
|
if self.args.log:
|
|
wandb.log(
|
|
{
|
|
"total_steps": self.total_steps,
|
|
"score": score,
|
|
"total loss": total_loss,
|
|
"actor loss": loss[0] * self.hyper_params["POLICY_UPDATE_FREQ"],
|
|
"critic1 loss": loss[1],
|
|
"critic2 loss": loss[2],
|
|
}
|
|
)
|
|
|
|
def train(self):
|
|
"""Train the agent."""
|
|
# logger
|
|
if self.args.log:
|
|
wandb.init()
|
|
wandb.config.update(self.hyper_params)
|
|
wandb.config.update(vars(self.args))
|
|
wandb.watch([self.actor, self.critic1, self.critic2], log="parameters")
|
|
|
|
for i_episode in range(1, self.args.episode_num + 1):
|
|
state = self.env.reset()
|
|
done = False
|
|
score = 0
|
|
loss_episode = list()
|
|
self.episode_steps = 0
|
|
|
|
while not done:
|
|
if self.args.render and i_episode >= self.args.render_after:
|
|
self.env.render()
|
|
|
|
action = self.select_action(state)
|
|
next_state, reward, done = self.step(action)
|
|
|
|
if len(self.memory) >= self.hyper_params["BATCH_SIZE"]:
|
|
experiences = self.memory.sample()
|
|
loss = self.update_model(experiences)
|
|
loss_episode.append(loss) # for logging
|
|
|
|
state = next_state
|
|
score += reward
|
|
|
|
# logging
|
|
if loss_episode:
|
|
avg_loss = np.vstack(loss_episode).mean(axis=0)
|
|
self.write_log(i_episode, avg_loss, score)
|
|
|
|
if i_episode % self.args.save_period == 0:
|
|
self.save_params(i_episode)
|
|
|
|
# termination
|
|
self.env.close()
|