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/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()