mirror of
https://github.com/wassname/kair_algorithms_draft.git
synced 2026-09-10 12:14:41 +08:00
Add overall setting and ddpg baseline (#1)
* Add overall CI settings * Add specific build dir to travis * Add before install/script condition to travis * Add ddpg baseline * Add wandb, remove algorithms except ddpg * Remove init file in script * Separate config file for ddpg * Remove unnecessary examples * Remove unnecessary args opt * Add pre-commit setting * Change pre-commit settings * Change travis-ci setting * Fix travis-ci issue * Modify argparse arguments, fix requirements * Change arguments order
This commit is contained in:
Whi Kwon
@@ -0,0 +1,139 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Abstract Agent used for all agents.
|
||||
|
||||
- Author: Curt Park
|
||||
- Contact: curt.park@medipixel.io
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Tuple
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class AbstractAgent(ABC):
|
||||
"""Abstract Agent used for all agents.
|
||||
|
||||
Attributes:
|
||||
env (gym.Env): openAI Gym environment with discrete action space
|
||||
args (argparse.Namespace): arguments including hyperparameters and training settings
|
||||
state_dim (int): dimension of state space
|
||||
action_dim (int): dimension of action space
|
||||
sha (str): sha code of current git commit
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, env: gym.Env, args: argparse.Namespace):
|
||||
"""Initialization.
|
||||
|
||||
Args:
|
||||
env (gym.Env): openAI Gym environment with discrete action space
|
||||
args (argparse.Namespace): arguments including hyperparameters and training settings
|
||||
|
||||
"""
|
||||
self.args = args
|
||||
self.env = NormalizedActions(env)
|
||||
if self.args.max_episode_steps > 0:
|
||||
env._max_episode_steps = self.args.max_episode_steps
|
||||
else:
|
||||
self.args.max_episode_steps = env._max_episode_steps
|
||||
|
||||
# for logging
|
||||
self.sha = (
|
||||
subprocess.check_output(["git", "rev-parse", "--short", "HEAD"])[:-1]
|
||||
.decode("ascii")
|
||||
.strip()
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def select_action(self, state: np.ndarray):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def step(self, action: torch.Tensor) -> Tuple[np.ndarray, np.float64, bool]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def update_model(self, *args):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def load_params(self, *args):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save_params(self, name: str, params: dict, n_episode: int):
|
||||
if not os.path.exists("./save"):
|
||||
os.mkdir("./save")
|
||||
|
||||
path = os.path.join(
|
||||
"./save/" + name + "_" + self.sha + "_ep_" + str(n_episode) + ".pt"
|
||||
)
|
||||
torch.save(params, path)
|
||||
|
||||
print("[INFO] Saved the model and optimizer to", path)
|
||||
|
||||
@abstractmethod
|
||||
def write_log(self, *args):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def train(self):
|
||||
pass
|
||||
|
||||
def test(self):
|
||||
"""Test the agent."""
|
||||
for i_episode in range(self.args.episode_num):
|
||||
state = self.env.reset()
|
||||
done = False
|
||||
score = 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)
|
||||
|
||||
state = next_state
|
||||
score += reward
|
||||
|
||||
print("[INFO] episode %d\ttotal score: %d" % (i_episode, score))
|
||||
|
||||
# termination
|
||||
self.env.close()
|
||||
|
||||
|
||||
class NormalizedActions(gym.ActionWrapper):
|
||||
"""Rescale and relocate the actions."""
|
||||
|
||||
def action(self, action: np.ndarray) -> np.ndarray:
|
||||
"""Change the range (-1, 1) to (low, high)."""
|
||||
low = self.action_space.low
|
||||
high = self.action_space.high
|
||||
|
||||
scale_factor = (high - low) / 2
|
||||
reloc_factor = high - scale_factor
|
||||
|
||||
action = action * scale_factor + reloc_factor
|
||||
action = np.clip(action, low, high)
|
||||
|
||||
return action
|
||||
|
||||
def reverse_action(self, action: np.ndarray) -> np.ndarray:
|
||||
"""Change the range (low, high) to (-1, 1)."""
|
||||
low = self.action_space.low
|
||||
high = self.action_space.high
|
||||
|
||||
scale_factor = (high - low) / 2
|
||||
reloc_factor = high - scale_factor
|
||||
|
||||
action = (action - reloc_factor) / scale_factor
|
||||
action = np.clip(action, -1.0, 1.0)
|
||||
|
||||
return action
|
||||
@@ -0,0 +1,72 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Replay buffer for baselines."""
|
||||
|
||||
import random
|
||||
from collections import deque
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class ReplayBuffer:
|
||||
"""Fixed-size buffer to store experience tuples.
|
||||
|
||||
Taken from Udacity deep-reinforcement-learning github repository:
|
||||
https://github.com/udacity/deep-reinforcement-learning/blob/master/
|
||||
ddpg-pendulum/ddpg_agent.py
|
||||
|
||||
Attributes:
|
||||
buffer (deque): deque of replay buffer
|
||||
batch_size (int): size of a batched sampled from replay buffer for training
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, buffer_size, batch_size, seed, demo=None):
|
||||
"""Initialize a ReplayBuffer object.
|
||||
|
||||
Args:
|
||||
buffer_size (int): size of replay buffer for experience
|
||||
batch_size (int): size of a batched sampled from replay buffer for training
|
||||
seed (int): random seed
|
||||
demo (deque) : demonstration deque
|
||||
|
||||
"""
|
||||
self.buffer = deque(maxlen=buffer_size) if not demo else demo
|
||||
|
||||
self.batch_size = batch_size
|
||||
random.seed(seed)
|
||||
|
||||
def add(self, state, action, reward, next_state, done):
|
||||
"""Add a new experience to memory."""
|
||||
self.buffer.append((state, action, reward, next_state, done))
|
||||
|
||||
def extend(self, transitions):
|
||||
"""Add experiences to memory."""
|
||||
self.buffer.extend(transitions)
|
||||
|
||||
def sample(self):
|
||||
"""Randomly sample a batch of experiences from memory."""
|
||||
experiences = random.sample(self.buffer, k=self.batch_size)
|
||||
|
||||
states, actions, rewards, next_states, dones = [], [], [], [], []
|
||||
|
||||
for e in experiences:
|
||||
states.append(np.expand_dims(e[0], axis=0))
|
||||
actions.append(e[1])
|
||||
rewards.append(e[2])
|
||||
next_states.append(np.expand_dims(e[3], axis=0))
|
||||
dones.append(e[4])
|
||||
|
||||
states = torch.from_numpy(np.vstack(states)).float().to(device)
|
||||
actions = torch.from_numpy(np.vstack(actions)).float().to(device)
|
||||
rewards = torch.from_numpy(np.vstack(rewards)).float().to(device)
|
||||
next_states = torch.from_numpy(np.vstack(next_states)).float().to(device)
|
||||
dones = torch.from_numpy(np.vstack(dones).astype(np.uint8)).float().to(device)
|
||||
|
||||
return (states, actions, rewards, next_states, dones)
|
||||
|
||||
def __len__(self):
|
||||
"""Return the current size of internal memory."""
|
||||
return len(self.buffer)
|
||||
@@ -0,0 +1,20 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Common util functions for all algorithms.
|
||||
|
||||
- Author: Curt Park
|
||||
- Contact: curt.park@medipixel.io
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def identity(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Return input without any change."""
|
||||
return x
|
||||
|
||||
|
||||
def soft_update(local: nn.Module, target: nn.Module, tau: float):
|
||||
"""Soft-update: target = tau*local + (1-tau)*target."""
|
||||
for t_param, l_param in zip(target.parameters(), local.parameters()):
|
||||
t_param.data.copy_(tau * l_param.data + (1.0 - tau) * t_param.data)
|
||||
@@ -0,0 +1,204 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""MLP module for model of algorithms
|
||||
|
||||
- Author: Kh Kim
|
||||
- Contact: kh.kim@medipixel.io
|
||||
"""
|
||||
|
||||
from typing import Callable, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.distributions import Normal
|
||||
|
||||
from algorithms.common.helper_functions import identity
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""Baseline of Multilayer perceptron.
|
||||
|
||||
Attributes:
|
||||
input_size (int): size of input
|
||||
output_size (int): size of output layer
|
||||
hidden_sizes (list): sizes of hidden layers
|
||||
hidden_activation (function): activation function of hidden layers
|
||||
output_activation (function): activation function of output layer
|
||||
hidden_layers (list): list containing linear layers
|
||||
use_output_layer (bool): whether or not to use the last layer
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
hidden_sizes: list,
|
||||
hidden_activation: Callable = F.relu,
|
||||
output_activation: Callable = identity,
|
||||
use_output_layer: bool = True,
|
||||
init_w: float = 3e-3,
|
||||
):
|
||||
"""Initialization.
|
||||
|
||||
Args:
|
||||
input_size (int): size of input
|
||||
output_size (int): size of output layer
|
||||
hidden_sizes (list): number of hidden layers
|
||||
hidden_activation (function): activation function of hidden layers
|
||||
output_activation (function): activation function of output layer
|
||||
use_output_layer (bool): whether or not to use the last layer
|
||||
init_w (float): weight initialization bound for the last layer
|
||||
|
||||
"""
|
||||
super(MLP, self).__init__()
|
||||
|
||||
self.hidden_sizes = hidden_sizes
|
||||
self.input_size = input_size
|
||||
self.output_size = output_size
|
||||
self.hidden_activation = hidden_activation
|
||||
self.output_activation = output_activation
|
||||
self.use_output_layer = use_output_layer
|
||||
|
||||
# set hidden layers
|
||||
self.hidden_layers: list = []
|
||||
in_size = self.input_size
|
||||
for i, next_size in enumerate(hidden_sizes):
|
||||
fc = nn.Linear(in_size, next_size)
|
||||
in_size = next_size
|
||||
self.__setattr__("hidden_fc{}".format(i), fc)
|
||||
self.hidden_layers.append(fc)
|
||||
|
||||
# set output layers
|
||||
if self.use_output_layer:
|
||||
self.output_layer = nn.Linear(in_size, output_size)
|
||||
self.output_layer.weight.data.uniform_(-init_w, init_w)
|
||||
self.output_layer.bias.data.uniform_(-init_w, init_w)
|
||||
|
||||
def get_last_activation(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Get the activation of the last hidden layer."""
|
||||
for hidden_layer in self.hidden_layers:
|
||||
x = self.hidden_activation(hidden_layer(x))
|
||||
return x
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Forward method implementation."""
|
||||
assert self.use_output_layer
|
||||
|
||||
x = self.get_last_activation(x)
|
||||
|
||||
output = self.output_layer(x)
|
||||
output = self.output_activation(output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class GaussianDist(MLP):
|
||||
"""Multilayer perceptron with Gaussian distribution output.
|
||||
|
||||
Attributes:
|
||||
mu_activation (function): bounding function for mean
|
||||
log_std_clamping (bool): whether or not to clamp log std
|
||||
log_std_min (float): lower bound of log std
|
||||
log_std_max (float): upper bound of log std
|
||||
mu_layer (nn.Linear): output layer for mean
|
||||
log_std_layer (nn.Linear): output layer for log std
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
hidden_sizes: list,
|
||||
hidden_activation: Callable = F.relu,
|
||||
mu_activation: Callable = torch.tanh,
|
||||
log_std_min: float = -20,
|
||||
log_std_max: float = 2,
|
||||
init_w: float = 3e-3,
|
||||
):
|
||||
"""Initialization.
|
||||
|
||||
"""
|
||||
super(GaussianDist, self).__init__(
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
hidden_sizes=hidden_sizes,
|
||||
hidden_activation=hidden_activation,
|
||||
use_output_layer=False,
|
||||
)
|
||||
|
||||
self.mu_activation = mu_activation
|
||||
self.log_std_min = log_std_min
|
||||
self.log_std_max = log_std_max
|
||||
in_size = hidden_sizes[-1]
|
||||
|
||||
# set log_std layer
|
||||
self.log_std_layer = nn.Linear(in_size, output_size)
|
||||
self.log_std_layer.weight.data.uniform_(-init_w, init_w)
|
||||
self.log_std_layer.bias.data.uniform_(-init_w, init_w)
|
||||
|
||||
# set mean layer
|
||||
self.mu_layer = nn.Linear(in_size, output_size)
|
||||
self.mu_layer.weight.data.uniform_(-init_w, init_w)
|
||||
self.mu_layer.bias.data.uniform_(-init_w, init_w)
|
||||
|
||||
def get_dist_params(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
|
||||
"""Return gausian distribution parameters."""
|
||||
hidden = super(GaussianDist, self).get_last_activation(x)
|
||||
|
||||
# get mean
|
||||
mu = self.mu_activation(self.mu_layer(hidden))
|
||||
|
||||
# get std
|
||||
log_std = torch.clamp(
|
||||
self.log_std_layer(hidden), self.log_std_min, self.log_std_max
|
||||
)
|
||||
std = torch.exp(log_std)
|
||||
|
||||
return mu, log_std, std
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
|
||||
"""Forward method implementation."""
|
||||
mu, _, std = self.get_dist_params(x)
|
||||
|
||||
# get normal distribution and action
|
||||
dist = Normal(mu, std)
|
||||
action = dist.sample()
|
||||
|
||||
return action, dist
|
||||
|
||||
|
||||
class GaussianDistParams(GaussianDist):
|
||||
"""Multilayer perceptron with Gaussian distribution params output."""
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
|
||||
"""Forward method implementation."""
|
||||
mu, log_std, std = super(GaussianDistParams, self).get_dist_params(x)
|
||||
|
||||
return mu, log_std, std
|
||||
|
||||
|
||||
class TanhGaussianDistParams(GaussianDist):
|
||||
"""Multilayer perceptron with Gaussian distribution output."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
"""Initialization."""
|
||||
super(TanhGaussianDistParams, self).__init__(**kwargs, mu_activation=identity)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, epsilon: float = 1e-6
|
||||
) -> Tuple[torch.Tensor, ...]:
|
||||
"""Forward method implementation."""
|
||||
mu, _, std = super(TanhGaussianDistParams, self).get_dist_params(x)
|
||||
|
||||
# sampling actions
|
||||
dist = Normal(mu, std)
|
||||
z = dist.rsample()
|
||||
|
||||
# normalize action and log_prob
|
||||
# see appendix C of 'https://arxiv.org/pdf/1812.05905.pdf'
|
||||
action = torch.tanh(z)
|
||||
log_prob = dist.log_prob(z) - torch.log(1 - action.pow(2) + epsilon)
|
||||
log_prob = log_prob.sum(-1, keepdim=True)
|
||||
|
||||
return action, log_prob, z, mu, std
|
||||
@@ -0,0 +1,39 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Noise classes for baselines."""
|
||||
|
||||
import copy
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class OUNoise:
|
||||
"""Ornstein-Uhlenbeck process.
|
||||
|
||||
Taken from Udacity deep-reinforcement-learning github repository:
|
||||
https://github.com/udacity/deep-reinforcement-learning/blob/master/
|
||||
ddpg-pendulum/ddpg_agent.py
|
||||
"""
|
||||
|
||||
def __init__(self, size, seed, mu=0.0, theta=0.15, sigma=0.2):
|
||||
"""Initialize parameters and noise process."""
|
||||
self.state = np.float64(0.0)
|
||||
self.mu = mu * np.ones(size)
|
||||
self.theta = theta
|
||||
self.sigma = sigma
|
||||
self.reset()
|
||||
|
||||
random.seed(seed)
|
||||
|
||||
def reset(self):
|
||||
"""Reset the internal state (= noise) to mean (mu)."""
|
||||
self.state = copy.copy(self.mu)
|
||||
|
||||
def sample(self):
|
||||
"""Update internal state and return it as a noise sample."""
|
||||
x = self.state
|
||||
dx = self.theta * (self.mu - x) + self.sigma * np.array(
|
||||
[random.random() for _ in range(len(x))]
|
||||
)
|
||||
self.state = x + dx
|
||||
return self.state
|
||||
@@ -0,0 +1,226 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""DDPG agent for episodic tasks in OpenAI Gym.
|
||||
|
||||
- Author: Curt Park
|
||||
- Contact: curt.park@medipixel.io
|
||||
- Paper: https://arxiv.org/pdf/1509.02971.pdf
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from typing import Tuple
|
||||
|
||||
import gym
|
||||
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
|
||||
from algorithms.common.noise import OUNoise
|
||||
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class Agent(AbstractAgent):
|
||||
"""ActorCritic interacting with environment.
|
||||
|
||||
Attributes:
|
||||
memory (ReplayBuffer): replay memory
|
||||
noise (OUNoise): random noise for exploration
|
||||
hyper_params (dict): hyper-parameters
|
||||
actor (nn.Module): actor model to select actions
|
||||
actor_target (nn.Module): target actor model to select actions
|
||||
critic (nn.Module): critic model to predict state values
|
||||
critic_target (nn.Module): target critic model to predict state values
|
||||
actor_optimizer (Optimizer): optimizer for training actor
|
||||
critic_optimizer (Optimizer): optimizer for training critic
|
||||
curr_state (np.ndarray): temporary storage of the current state
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
env: gym.Env,
|
||||
args: argparse.Namespace,
|
||||
hyper_params: dict,
|
||||
models: tuple,
|
||||
optims: tuple,
|
||||
noise: OUNoise,
|
||||
):
|
||||
"""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 critic
|
||||
optims (tuple): optimizers for actor and critic
|
||||
noise (OUNoise): random noise for exploration
|
||||
|
||||
"""
|
||||
AbstractAgent.__init__(self, env, args)
|
||||
|
||||
self.actor, self.actor_target, self.critic, self.critic_target = models
|
||||
self.actor_optimizer, self.critic_optimizer = optims
|
||||
self.hyper_params = hyper_params
|
||||
self.curr_state = np.zeros((1,))
|
||||
self.noise = noise
|
||||
|
||||
# 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)
|
||||
|
||||
# replay memory
|
||||
self.memory = ReplayBuffer(
|
||||
hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"], self.args.seed
|
||||
)
|
||||
|
||||
def select_action(self, state: np.ndarray) -> torch.Tensor:
|
||||
"""Select an action from the input space."""
|
||||
self.curr_state = state
|
||||
|
||||
state = torch.FloatTensor(state).to(device)
|
||||
selected_action = self.actor(state)
|
||||
selected_action += torch.FloatTensor(self.noise.sample()).to(device)
|
||||
|
||||
selected_action = torch.clamp(selected_action, -1.0, 1.0)
|
||||
|
||||
return selected_action
|
||||
|
||||
def step(self, action: torch.Tensor) -> Tuple[np.ndarray, np.float64, bool]:
|
||||
"""Take an action and return the response of the env."""
|
||||
action = action.detach().cpu().numpy()
|
||||
next_state, reward, done, _ = self.env.step(action)
|
||||
|
||||
self.memory.add(self.curr_state, action, reward, next_state, done)
|
||||
|
||||
return next_state, reward, done
|
||||
|
||||
def update_model(
|
||||
self,
|
||||
experiences: Tuple[
|
||||
torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
|
||||
],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""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
|
||||
next_actions = self.actor_target(next_states)
|
||||
next_values = self.critic_target(torch.cat((next_states, next_actions), dim=-1))
|
||||
curr_returns = rewards + self.hyper_params["GAMMA"] * next_values * masks
|
||||
curr_returns = curr_returns.to(device)
|
||||
|
||||
# train critic
|
||||
values = self.critic(torch.cat((states, actions), dim=-1))
|
||||
critic_loss = F.mse_loss(values, curr_returns)
|
||||
self.critic_optimizer.zero_grad()
|
||||
critic_loss.backward()
|
||||
self.critic_optimizer.step()
|
||||
|
||||
# train actor
|
||||
actions = self.actor(states)
|
||||
actor_loss = -self.critic(torch.cat((states, actions), dim=-1)).mean()
|
||||
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)
|
||||
|
||||
return actor_loss.data, critic_loss.data
|
||||
|
||||
def load_params(self, path: str):
|
||||
"""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.critic.load_state_dict(params["critic_state_dict"])
|
||||
self.critic_target.load_state_dict(params["critic_target_state_dict"])
|
||||
self.actor_optimizer.load_state_dict(params["actor_optim_state_dict"])
|
||||
self.critic_optimizer.load_state_dict(params["critic_optim_state_dict"])
|
||||
print("[INFO] loaded the model and optimizer from", path)
|
||||
|
||||
def save_params(self, n_episode: int):
|
||||
"""Save model and optimizer parameters."""
|
||||
params = {
|
||||
"actor_state_dict": self.actor.state_dict(),
|
||||
"actor_target_state_dict": self.actor_target.state_dict(),
|
||||
"critic_state_dict": self.critic.state_dict(),
|
||||
"critic_target_state_dict": self.critic_target.state_dict(),
|
||||
"actor_optim_state_dict": self.actor_optimizer.state_dict(),
|
||||
"critic_optim_state_dict": self.critic_optimizer.state_dict(),
|
||||
}
|
||||
|
||||
AbstractAgent.save_params(self, self.args.algo, params, n_episode)
|
||||
|
||||
def write_log(self, i: int, loss: np.ndarray, score: int):
|
||||
"""Write log about loss and score"""
|
||||
total_loss = loss.sum()
|
||||
|
||||
print(
|
||||
"[INFO] episode %d total score: %d, total loss: %f\n"
|
||||
"actor_loss: %.3f critic_loss: %.3f\n"
|
||||
% (i, score, total_loss, loss[0], loss[1]) # actor loss # critic loss
|
||||
)
|
||||
|
||||
if self.args.log:
|
||||
wandb.log(
|
||||
{
|
||||
"score": score,
|
||||
"total loss": total_loss,
|
||||
"actor loss": loss[0],
|
||||
"critic loss": loss[1],
|
||||
}
|
||||
)
|
||||
|
||||
def train(self):
|
||||
"""Train the agent."""
|
||||
# logger
|
||||
if self.args.log:
|
||||
wandb.init()
|
||||
wandb.config.update(self.hyper_params)
|
||||
wandb.watch([self.actor, self.critic], log="parameters")
|
||||
|
||||
for i_episode in range(1, self.args.episode_num + 1):
|
||||
state = self.env.reset()
|
||||
done = False
|
||||
score = 0
|
||||
loss_episode = list()
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user