Add per (#8)

* Add per and modify etc

* Replace pre-commit-config.yaml and add pre-commit hook in .git

* Modify .gitignore

* Modify .gitignore

* Modify buffer and code

* Modify replay buffer and per

* Modify .gitignore

scripts/algorithms/common/abstract/agent.py

@@ -20,10 +20,9 @@ class AbstractAgent(ABC):
     """Abstract Agent used for all agents.
 
     Attributes:
-        env (gym.Env): openAI Gym environment with discrete action space
+        env (gym.Env): openAI Gym environment
         args (argparse.Namespace): arguments including hyperparameters and training settings
-        state_dim (int): dimension of state space
-        action_dim (int): dimension of action space
+        env_name (str) : gym env name for logging
         sha (str): sha code of current git commit
 
     """
@@ -32,7 +31,7 @@ class AbstractAgent(ABC):
         """Initialization.
 
         Args:
-            env (gym.Env): openAI Gym environment with discrete action space
+            env (gym.Env): openAI Gym environment
             args (argparse.Namespace): arguments including hyperparameters and training settings
 
         """
@@ -44,6 +43,7 @@ class AbstractAgent(ABC):
             self.args.max_episode_steps = env._max_episode_steps
 
         # for logging
+        self.env_name = str(self.env.env).split("<")[2].replace(">>", "")
         self.sha = (
             subprocess.check_output(["git", "rev-parse", "--short", "HEAD"])[:-1]
             .decode("ascii")
@@ -67,13 +67,13 @@ class AbstractAgent(ABC):
         pass
 
     @abstractmethod
-    def save_params(self, name: str, params: dict, n_episode: int):
+    def save_params(self, 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"
-        )
+        save_name = self.env_name + "_" + self.args.algo + "_" + self.sha
+
+        path = os.path.join("./save/" + save_name + "_ep_" + str(n_episode) + ".pt")
         torch.save(params, path)
 
         print("[INFO] Saved the model and optimizer to", path)
@@ -92,6 +92,7 @@ class AbstractAgent(ABC):
             state = self.env.reset()
             done = False
             score = 0
+            step = 0
 
             while not done:
                 if self.args.render and i_episode >= self.args.render_after:
@@ -102,8 +103,12 @@ class AbstractAgent(ABC):
 
                 state = next_state
                 score += reward
+                step += 1
 
-            print("[INFO] episode %d\ttotal score: %d" % (i_episode, score))
+            print(
+                "[INFO] episode %d\tstep: %d\ttotal score: %d"
+                % (i_episode, step, score)
+            )
 
         # termination
         self.env.close()

scripts/algorithms/common/buffer/priortized_replay_buffer.py

@@ -0,0 +1,151 @@
+# -*- coding: utf-8 -*-
+"""Prioritized Replay buffer for baselines.
+
+- Author: Kh Kim
+- Contact: kh.kim@medipixel.io
+- Paper: https://arxiv.org/pdf/1511.05952.pdf
+         https://arxiv.org/pdf/1707.08817.pdf
+"""
+
+import random
+from typing import Tuple
+
+import numpy as np
+import torch
+
+from algorithms.common.buffer.replay_buffer import ReplayBuffer
+from algorithms.common.buffer.segment_tree import MinSegmentTree, SumSegmentTree
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+class PrioritizedReplayBuffer(ReplayBuffer):
+    """Create Prioritized Replay buffer.
+
+    Taken from OpenAI baselines github repository:
+    https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py
+
+    Attributes:
+        buffer_size (int): size of replay buffer for experience
+        alpha (float): alpha parameter for prioritized replay buffer
+        tree_idx (int): next index of tree
+        sum_tree (SumSegmentTree): sum tree for prior
+        min_tree (MinSegmentTree): min tree for min prior to get max weight
+        init_priority (float): lower bound of priority
+
+        """
+
+    def __init__(
+        self, buffer_size: int, batch_size: int, demo: list = None, alpha: float = 0.6
+    ):
+        """Initialization.
+
+        Args:
+            buffer_size (int): size of replay buffer for experience
+            batch_size (int): size of a batched sampled from replay buffer for training
+            demo (list): demonstration
+            alpha (float): alpha parameter for prioritized replay buffer
+
+        """
+        super(PrioritizedReplayBuffer, self).__init__(buffer_size, batch_size, demo)
+        assert alpha >= 0
+        self.buffer_size = buffer_size
+        self.alpha = alpha
+        self.tree_idx = 0
+
+        # capacity must be positive and a power of 2.
+        tree_capacity = 1
+        while tree_capacity < self.buffer_size:
+            tree_capacity *= 2
+
+        self.sum_tree = SumSegmentTree(tree_capacity)
+        self.min_tree = MinSegmentTree(tree_capacity)
+        self.init_priority = 1.0
+
+        # for init priority of demo
+        if demo:
+            for _ in range(len(demo)):
+                self.sum_tree[self.tree_idx] = self.init_priority ** self.alpha
+                self.min_tree[self.tree_idx] = self.init_priority ** self.alpha
+                self.tree_idx += 1
+
+    def add(
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        reward: np.float64,
+        next_state: np.ndarray,
+        done: bool,
+    ):
+        """Add experience and priority."""
+        idx = self.tree_idx
+        self.tree_idx = (self.tree_idx + 1) % self.buffer_size
+        super().add(state, action, reward, next_state, done)
+
+        self.sum_tree[idx] = self.init_priority ** self.alpha
+        self.min_tree[idx] = self.init_priority ** self.alpha
+
+    def extend(self, transitions: list):
+        """Add experiences to memory."""
+        raise NotImplementedError
+
+    def _sample_proportional(self, batch_size: int) -> list:
+        """Sample indices based on proportional."""
+        indices = []
+        p_total = self.sum_tree.sum(0, len(self.buffer) - 1)
+        segment = p_total / batch_size
+        for i in range(batch_size):
+            a = segment * i
+            b = segment * (i + 1)
+            upperbound = random.uniform(a, b)
+            idx = self.sum_tree.retrieve(upperbound)
+            indices.append(idx)
+        return indices
+
+    def sample(self, beta: float = 0.4) -> Tuple[torch.Tensor, ...]:
+        """Sample a batch of experiences."""
+        assert beta > 0
+
+        indices = self._sample_proportional(self.batch_size)
+        states, actions, rewards, next_states, dones, weights = [], [], [], [], [], []
+
+        # get max weight
+        p_min = self.min_tree.min() / self.sum_tree.sum()
+        max_weight = (p_min * len(self.buffer)) ** (-beta)
+
+        for i in indices:
+            s, a, r, n_s, d = self.buffer[i]
+            states.append(np.array(s, copy=False))
+            actions.append(np.array(a, copy=False))
+            rewards.append(np.array(r, copy=False))
+            next_states.append(np.array(n_s, copy=False))
+            dones.append(np.array(float(d), copy=False))
+
+            # calculate weights
+            p_sample = self.sum_tree[i] / self.sum_tree.sum()
+            weight = (p_sample * len(self.buffer)) ** (-beta)
+            weights.append(weight / max_weight)
+
+        states = torch.FloatTensor(np.array(states)).to(device)
+        actions = torch.FloatTensor(np.array(actions)).to(device)
+        rewards = torch.FloatTensor(np.array(rewards).reshape(-1, 1)).to(device)
+        next_states = torch.FloatTensor(np.array(next_states)).to(device)
+        dones = torch.FloatTensor(np.array(dones).reshape(-1, 1)).to(device)
+        weights = torch.FloatTensor(np.array(weights).reshape(-1, 1)).to(device)
+
+        experiences = (states, actions, rewards, next_states, dones, weights, indices)
+
+        return experiences
+
+    def update_priorities(self, indices: list, priorities: np.ndarray):
+        """Update priorities of sampled transitions."""
+        assert len(indices) == len(priorities)
+
+        for idx, priority in zip(indices, priorities):
+            assert priority > 0
+            assert 0 <= idx < len(self.buffer)
+
+            self.sum_tree[idx] = priority ** self.alpha
+            self.min_tree[idx] = priority ** self.alpha
+
+            self.init_priority = max(self.init_priority, priority)

scripts/algorithms/common/buffer/replay_buffer.py

@@ -1,8 +1,7 @@
 # -*- coding: utf-8 -*-
 """Replay buffer for baselines."""
 
-import random
-from collections import deque
+from typing import Tuple
 
 import numpy as np
 import torch
@@ -18,55 +17,70 @@ class ReplayBuffer:
     ddpg-pendulum/ddpg_agent.py
 
     Attributes:
-        buffer (deque): deque of replay buffer
+        buffer (list): list 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):
+    def __init__(self, buffer_size: int, batch_size: int, demo: list = 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
+            demo (list) : demonstration list
 
         """
-        self.buffer = deque(maxlen=buffer_size) if not demo else demo
-
+        self.buffer = list() if not demo else demo
+        self.buffer_size = buffer_size
         self.batch_size = batch_size
-        random.seed(seed)
+        self.idx = 0
 
-    def add(self, state, action, reward, next_state, done):
+    def add(
+        self,
+        state: np.ndarray,
+        action: np.ndarray,
+        reward: np.float64,
+        next_state: np.ndarray,
+        done: bool,
+    ):
         """Add a new experience to memory."""
-        self.buffer.append((state, action, reward, next_state, done))
+        data = (state, action, reward, next_state, done)
 
-    def extend(self, transitions):
+        if len(self.buffer) == self.buffer_size:
+            self.buffer[self.idx] = data
+            self.idx = (self.idx + 1) % self.buffer_size
+        else:
+            self.buffer.append(data)
+
+    def extend(self, transitions: list):
         """Add experiences to memory."""
         self.buffer.extend(transitions)
 
-    def sample(self):
+    def sample(
+        self
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """Randomly sample a batch of experiences from memory."""
-        experiences = random.sample(self.buffer, k=self.batch_size)
+        idxs = np.random.randint(0, len(self.buffer), size=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])
+        for i in idxs:
+            s, a, r, n_s, d = self.buffer[i]
+            states.append(np.array(s, copy=False))
+            actions.append(np.array(a, copy=False))
+            rewards.append(np.array(r, copy=False))
+            next_states.append(np.array(n_s, copy=False))
+            dones.append(np.array(float(d), copy=False))
 
-        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)
+        states = torch.FloatTensor(np.array(states)).to(device)
+        actions = torch.FloatTensor(np.array(actions)).to(device)
+        rewards = torch.FloatTensor(np.array(rewards).reshape(-1, 1)).to(device)
+        next_states = torch.FloatTensor(np.array(next_states)).to(device)
+        dones = torch.FloatTensor(np.array(dones).reshape(-1, 1)).to(device)
 
-        return (states, actions, rewards, next_states, dones)
+        return states, actions, rewards, next_states, dones
 
-    def __len__(self):
+    def __len__(self) -> int:
         """Return the current size of internal memory."""
         return len(self.buffer)

scripts/algorithms/common/buffer/segment_tree.py

@@ -0,0 +1,141 @@
+# -*- coding: utf-8 -*-
+"""Segment tree for Proirtized Replay Buffer."""
+
+import operator
+from typing import Callable
+
+
+class SegmentTree:
+    """ Create SegmentTree.
+
+    Taken from OpenAI baselines github repository:
+    https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
+
+    Attributes:
+        capacity (int)
+        tree (list)
+        operation (function)
+
+    """
+
+    def __init__(self, capacity: int, operation: Callable, init_value: float):
+        """Initialization.
+
+        Args:
+            capacity (int)
+            operation (function)
+            init_value (float)
+
+        """
+        assert (
+            capacity > 0 and capacity & (capacity - 1) == 0
+        ), "capacity must be positive and a power of 2."
+        self.capacity = capacity
+        self.tree = [init_value for _ in range(2 * capacity)]
+        self.operation = operation
+
+    def _operate_helper(
+        self, start: int, end: int, node: int, node_start: int, node_end: int
+    ) -> float:
+        """Returns result of operation in segment."""
+        if start == node_start and end == node_end:
+            return self.tree[node]
+        mid = (node_start + node_end) // 2
+        if end <= mid:
+            return self._operate_helper(start, end, 2 * node, node_start, mid)
+        else:
+            if mid + 1 <= start:
+                return self._operate_helper(start, end, 2 * node + 1, mid + 1, node_end)
+            else:
+                return self.operation(
+                    self._operate_helper(start, mid, 2 * node, node_start, mid),
+                    self._operate_helper(mid + 1, end, 2 * node + 1, mid + 1, node_end),
+                )
+
+    def operate(self, start: int = 0, end: int = 0) -> float:
+        """Returns result of applying `self.operation`."""
+        if end <= 0:
+            end += self.capacity
+        end -= 1
+
+        return self._operate_helper(start, end, 1, 0, self.capacity - 1)
+
+    def __setitem__(self, idx: int, val: float):
+        """Set value in tree."""
+        idx += self.capacity
+        self.tree[idx] = val
+
+        idx //= 2
+        while idx >= 1:
+            self.tree[idx] = self.operation(self.tree[2 * idx], self.tree[2 * idx + 1])
+            idx //= 2
+
+    def __getitem__(self, idx: int) -> float:
+        """Get real value in leaf node of tree."""
+        assert 0 <= idx < self.capacity
+
+        return self.tree[self.capacity + idx]
+
+
+class SumSegmentTree(SegmentTree):
+    """ Create SumSegmentTree.
+
+    Taken from OpenAI baselines github repository:
+    https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
+
+    """
+
+    def __init__(self, capacity: int):
+        """Initialization.
+
+        Args:
+            capacity (int)
+
+        """
+        super(SumSegmentTree, self).__init__(
+            capacity=capacity, operation=operator.add, init_value=0.0
+        )
+
+    def sum(self, start: int = 0, end: int = 0) -> float:
+        """Returns arr[start] + ... + arr[end]."""
+        return super(SumSegmentTree, self).operate(start, end)
+
+    def retrieve(self, upperbound: float) -> int:
+        """Find the highest index `i` about upper bound in the tree"""
+        assert 0 <= upperbound <= self.sum() + 1e-5
+
+        idx = 1
+
+        while idx < self.capacity:  # while non-leaf
+            left = 2 * idx
+            right = left + 1
+            if self.tree[left] > upperbound:
+                idx = 2 * idx
+            else:
+                upperbound -= self.tree[left]
+                idx = right
+        return idx - self.capacity
+
+
+class MinSegmentTree(SegmentTree):
+    """ Create SegmentTree.
+
+    Taken from OpenAI baselines github repository:
+    https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
+
+    """
+
+    def __init__(self, capacity: int):
+        """Initialization.
+
+        Args:
+            capacity (int)
+
+        """
+        super(MinSegmentTree, self).__init__(
+            capacity=capacity, operation=min, init_value=float("inf")
+        )
+
+    def min(self, start: int = 0, end: int = 0) -> float:
+        """Returns min(arr[start], ...,  arr[end])."""
+        return super(MinSegmentTree, self).operate(start, end)

scripts/algorithms/common/helper_functions.py

@@ -5,6 +5,10 @@
 - Contact: curt.park@medipixel.io
 """
 
+import random
+
+import gym
+import numpy as np
 import torch
 import torch.nn as nn
 
@@ -18,3 +22,11 @@ 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)
+
+
+def set_random_seed(seed: int, env: gym.Env):
+    """Set random seed"""
+    env.seed(seed)
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)

scripts/algorithms/common/noise.py

@@ -15,7 +15,7 @@ class OUNoise:
     ddpg-pendulum/ddpg_agent.py
     """
 
-    def __init__(self, size, seed, mu=0.0, theta=0.15, sigma=0.2):
+    def __init__(self, size, 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)
@@ -23,8 +23,6 @@ class OUNoise:
         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)

scripts/algorithms/ddpg/agent.py

@@ -53,7 +53,7 @@ class Agent(AbstractAgent):
         """Initialization.
 
         Args:
-            env (gym.Env): openAI Gym environment with discrete action space
+            env (gym.Env): openAI Gym environment
             args (argparse.Namespace): arguments including hyperparameters and training settings
             hyper_params (dict): hyper-parameters
             models (tuple): models including actor and critic
@@ -75,7 +75,7 @@ class Agent(AbstractAgent):
 
         # replay memory
         self.memory = ReplayBuffer(
-            hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"], self.args.seed
+            hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"]
         )
 
     def select_action(self, state: np.ndarray) -> torch.Tensor:
@@ -84,9 +84,10 @@ class Agent(AbstractAgent):
 
         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)
+        if not self.args.test:
+            selected_action += torch.FloatTensor(self.noise.sample()).to(device)
+            selected_action = torch.clamp(selected_action, -1.0, 1.0)
 
         return selected_action
 
@@ -163,7 +164,7 @@ class Agent(AbstractAgent):
             "critic_optim_state_dict": self.critic_optimizer.state_dict(),
         }
 
-        AbstractAgent.save_params(self, self.args.algo, params, n_episode)
+        AbstractAgent.save_params(self, params, n_episode)
 
     def write_log(self, i: int, loss: np.ndarray, score: int):
         """Write log about loss and score"""

scripts/algorithms/per/ddpg_agent.py

@@ -0,0 +1,251 @@
+# -*- coding: utf-8 -*-
+"""DDPG agent with PER 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
+"""
+
+import argparse
+import os
+from typing import List, Tuple
+
+import gym
+import numpy as np
+import torch
+import wandb
+
+import algorithms.common.helper_functions as common_utils
+from algorithms.common.abstract.agent import AbstractAgent
+from algorithms.common.buffer.priortized_replay_buffer import PrioritizedReplayBuffer
+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 (PrioritizedReplayBuffer): replay memory
+        noise (OUNoise): random noise for exploration
+        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
+        hyper_params (dict): hyper-parameters
+        beta (float): beta parameter for prioritized replay buffer
+        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
+            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.beta = self.hyper_params["PER_BETA"]
+        self.memory = PrioritizedReplayBuffer(
+            self.hyper_params["BUFFER_SIZE"],
+            self.hyper_params["BATCH_SIZE"],
+            alpha=self.hyper_params["PER_ALPHA"],
+        )
+
+    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)
+
+        if not self.args.test:
+            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,
+            torch.Tensor,
+            List[int],
+        ],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Train the model after each episode."""
+
+        states, actions, rewards, next_states, dones, weights, indexes = 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).detach()
+
+        # train critic
+        values = self.critic(torch.cat((states, actions), dim=-1))
+        critic_loss = torch.mean((values - curr_returns).pow(2) * 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 in PER
+        new_priorities = (values - curr_returns).pow(2)
+        new_priorities = (
+            new_priorities.data.cpu().numpy() + self.hyper_params["PER_EPS"]
+        )
+        self.memory.update_priorities(indexes, new_priorities)
+
+        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, 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(self.beta)
+                    loss = self.update_model(experiences)
+                    loss_episode.append(loss)  # for logging
+
+                state = next_state
+                score += reward
+
+            # increase beta
+            fraction = min(float(i_episode) / self.args.max_episode_steps, 1.0)
+            self.beta = self.beta + fraction * (1.0 - self.beta)
+
+            # 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()