Add DDPGfD, TD3fD and SACfD (#22)

* Format repository

* Clone files from medipixel repo

* Fix DDPGfDAgent.update_model()

* Fix bug on _initialize()

* Add demo-path parameter and demo data

* Rename init_priority to _max_priority for PER

This makes PER and PERfD consistent.

* Make i_episode attribute of DDPGAgent

* Clone SAC code from medipixel repo

* Fix update_model() for SACfD

* Fix _initialize() for SACfD

* Add is_discrete attribute to AbstractAgent for SACfD

* Add i_episode attribute to SACAgent for SACfD

* Modularize DDPGAgent and SACAgent

* Modify hyperparameters for DDPGfD and SACfD

* Add NStepBuffer

* Add n-step to DDPGfD

* Add n-step to SACfD

* Add TD3fD without n-step

* Attempt to tune hyperparameters

* Remove discrete environment check in SAC

* Implement n-step on TD3fD

* Fix step function of TD3

No done check, and _add_transition_to_memory was not called.

* Fix actor loss calculation for TD3fD

* Attempt to tune hyperparameters

* Print both critic losses

* Fix typo bug

* Attempt to tune hyperparameters

* Fix bug in n-step demo retrieval

* Fix bug in n-step transition addition

scripts/algorithms/common/abstract/agent.py

@@ -37,6 +37,7 @@ class AbstractAgent(ABC):
         """
         self.args = args
         self.env = NormalizedActions(env)
+
         if self.args.max_episode_steps > 0:
             env._max_episode_steps = self.args.max_episode_steps
         else:

scripts/algorithms/common/buffer/priortized_replay_buffer.py

@@ -31,13 +31,11 @@ class PrioritizedReplayBuffer(ReplayBuffer):
         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
+        _max_priority (float): max priority
 
         """
 
-    def __init__(
-        self, buffer_size: int, batch_size: int, alpha: float = 0.6
-    ):
+    def __init__(self, buffer_size: int, batch_size: int, alpha: float = 0.6):
         """Initialization.
 
         Args:
@@ -59,7 +57,7 @@ class PrioritizedReplayBuffer(ReplayBuffer):
 
         self.sum_tree = SumSegmentTree(tree_capacity)
         self.min_tree = MinSegmentTree(tree_capacity)
-        self.init_priority = 1.0
+        self._max_priority = 1.0
 
     def add(
         self,
@@ -74,8 +72,8 @@ class PrioritizedReplayBuffer(ReplayBuffer):
         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
+        self.sum_tree[idx] = self._max_priority ** self.alpha
+        self.min_tree[idx] = self._max_priority ** self.alpha
 
     def extend(self, transitions: list):
         """Add experiences to memory."""
@@ -140,4 +138,143 @@ class PrioritizedReplayBuffer(ReplayBuffer):
             self.sum_tree[idx] = priority ** self.alpha
             self.min_tree[idx] = priority ** self.alpha
 
-            self.init_priority = max(self.init_priority, priority)
+            self._max_priority = max(self._max_priority, priority)
+
+
+class PrioritizedReplayBufferfD(PrioritizedReplayBuffer):
+    """Create Prioritized Replay buffer with demo.
+    Taken from OpenAI baselines github repository:
+    https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py
+    Attributes:
+        demo (list): list of demo replay buffer
+        buffer_size (int): size of replay buffer for experience
+        demo_size (int): size of replay buffer for demonstration
+        total_size (int): sum of demo size and number of samples of experience
+        epsilon_d (float) : epsilon_d parameter to update priority using demo
+        """
+
+    def __init__(
+        self,
+        buffer_size: int,
+        batch_size: int,
+        demo: list,
+        alpha: float = 0.6,
+        epsilon_d: float = 1.0,
+    ):
+        """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
+            epsilon_d (float) : epsilon_d parameter to update priority using demo
+        """
+        super(PrioritizedReplayBufferfD, self).__init__(buffer_size, batch_size, alpha)
+        self.demo = demo
+        self.demo_size = len(demo)
+        self.total_size = self.demo_size + len(self.buffer)
+        self.epsilon_d = epsilon_d
+
+        # for init priority of demo
+        for _ in range(self.demo_size):
+            self.sum_tree[self.tree_idx] = self._max_priority ** self.alpha
+            self.min_tree[self.tree_idx] = self._max_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
+        # buffer is full
+        if (self.tree_idx + 1) % (self.buffer_size + self.demo_size) == 0:
+            self.tree_idx = self.demo_size
+        else:
+            self.tree_idx = self.tree_idx + 1
+        super().add(state, action, reward, next_state, done)
+
+        self.sum_tree[idx] = self._max_priority ** self.alpha
+        self.min_tree[idx] = self._max_priority ** self.alpha
+
+        # update current total size
+        self.total_size = self.demo_size + len(self.buffer)
+
+    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, eps_d = [], []
+
+        # get max weight
+        p_min = self.min_tree.min() / self.sum_tree.sum()
+        max_weight = (p_min * self.total_size) ** (-beta)
+
+        for i in indices:
+            # sample from buffer
+            if i < self.demo_size:
+                s, a, r, n_s, d = self.demo[i]
+                eps_d.append(self.epsilon_d)
+            else:
+                s, a, r, n_s, d = self.buffer[i - self.demo_size]
+                eps_d.append(0.0)
+
+            # append transition info
+            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 * self.total_size) ** (-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)
+        eps_d = np.array(eps_d)
+
+        if torch.cuda.is_available():
+            states_ = states_.cuda(non_blocking=True)
+            actions_ = actions_.cuda(non_blocking=True)
+            rewards_ = rewards_.cuda(non_blocking=True)
+            next_states_ = next_states_.cuda(non_blocking=True)
+            dones_ = dones_.cuda(non_blocking=True)
+            weights_ = weights_.cuda(non_blocking=True)
+
+        experiences = (
+            states_,
+            actions_,
+            rewards_,
+            next_states_,
+            dones_,
+            weights_,
+            indices,
+            eps_d,
+        )
+
+        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 < self.total_size
+
+            self.sum_tree[idx] = priority ** self.alpha
+            self.min_tree[idx] = priority ** self.alpha
+
+            self._max_priority = max(self._max_priority, priority)

scripts/algorithms/common/buffer/replay_buffer.py

@@ -1,11 +1,14 @@
 # -*- coding: utf-8 -*-
 """Replay buffer for baselines."""
 
-from typing import Tuple
+from collections import deque
+from typing import Any, Deque, List, Tuple
 
 import numpy as np
 import torch
 
+from algorithms.common.helper_functions import get_n_step_info
+
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 
@@ -84,3 +87,88 @@ class ReplayBuffer:
     def __len__(self) -> int:
         """Return the current size of internal memory."""
         return len(self.buffer)
+
+
+class NStepTransitionBuffer:
+    """Fixed-size buffer to store experience tuples.
+
+    Attributes:
+        buffer (list): list of replay buffer
+        buffer_size (int): buffer size not storing demos
+        demo_size (int): size of a demo to permanently store in the buffer
+        cursor (int): position to store next transition coming in
+
+    """
+
+    def __init__(self, buffer_size: int, n_step: int, gamma: float, demo: list = None):
+        """Initialize a ReplayBuffer object.
+
+        Args:
+            buffer_size (int): size of replay buffer for experience
+            demo (list): demonstration transitions
+
+        """
+        assert buffer_size > 0
+
+        self.n_step_buffer: Deque = deque(maxlen=n_step)
+        self.buffer_size = buffer_size
+        self.buffer: list = list()
+        self.n_step = n_step
+        self.gamma = gamma
+        self.demo_size = 0
+        self.cursor = 0
+
+        # if demo exists
+        if demo:
+            self.demo_size = len(demo)
+            self.buffer.extend(demo)
+
+        self.buffer.extend([None] * self.buffer_size)
+
+    def add(self, transition: Tuple[np.ndarray, ...]) -> Tuple[Any, ...]:
+        """Add a new transition to memory."""
+        self.n_step_buffer.append(transition)
+
+        # single step transition is not ready
+        if len(self.n_step_buffer) < self.n_step:
+            return ()
+
+        # add a multi step transition
+        reward, next_state, done = get_n_step_info(self.n_step_buffer, self.gamma)
+        curr_state, action = self.n_step_buffer[0][:2]
+        new_transition = (curr_state, action, reward, next_state, done)
+
+        # insert the new transition to buffer
+        idx = self.demo_size + self.cursor
+        self.buffer[idx] = new_transition
+        self.cursor = (self.cursor + 1) % self.buffer_size
+
+        # return a single step transition to insert to replay buffer
+        return self.n_step_buffer[0]
+
+    def sample(self, indices: List[int]) -> Tuple[torch.Tensor, ...]:
+        """Randomly sample a batch of experiences from memory."""
+        states, actions, rewards, next_states, dones = [], [], [], [], []
+
+        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))
+
+        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)
+
+        if torch.cuda.is_available():
+            states_ = states_.cuda(non_blocking=True)
+            actions_ = actions_.cuda(non_blocking=True)
+            rewards_ = rewards_.cuda(non_blocking=True)
+            next_states_ = next_states_.cuda(non_blocking=True)
+            dones_ = dones_.cuda(non_blocking=True)
+
+        return states_, actions_, rewards_, next_states_, dones_

scripts/algorithms/common/helper_functions.py

@@ -6,6 +6,8 @@
 """
 
 import random
+from collections import deque
+from typing import Deque, List, Tuple
 
 import gym
 import numpy as np
@@ -32,3 +34,46 @@ def set_random_seed(seed: int, env: gym.Env):
     torch.manual_seed(seed)
     np.random.seed(seed)
     random.seed(seed)
+
+
+def get_n_step_info_from_demo(
+    demo: List, n_step: int, gamma: float
+) -> Tuple[List, List]:
+    """Return 1 step and n step demos."""
+    assert demo
+    assert n_step > 1
+
+    demos_1_step = list()
+    demos_n_step = list()
+    n_step_buffer: Deque = deque(maxlen=n_step)
+
+    for transition in demo:
+        n_step_buffer.append(transition)
+
+        if len(n_step_buffer) == n_step:
+            # add a single step transition
+            demos_1_step.append(n_step_buffer[0])
+
+            # add a multi step transition
+            curr_state, action = n_step_buffer[0][:2]
+            reward, next_state, done = get_n_step_info(n_step_buffer, gamma)
+            transition = (curr_state, action, reward, next_state, done)
+            demos_n_step.append(transition)
+
+    return demos_1_step, demos_n_step
+
+
+def get_n_step_info(
+    n_step_buffer: Deque, gamma: float
+) -> Tuple[np.int64, np.ndarray, bool]:
+    """Return n step reward, next state, and done."""
+    # info of the last transition
+    reward, next_state, done = n_step_buffer[-1][-3:]
+
+    for transition in reversed(list(n_step_buffer)[:-1]):
+        r, n_s, d = transition[-3:]
+
+        reward = r + gamma * reward * (1 - d)
+        next_state, done = (n_s, d) if d else (next_state, done)
+
+    return reward, next_state, done

scripts/algorithms/ddpg/agent.py

@@ -40,6 +40,7 @@ class Agent(AbstractAgent):
         curr_state (np.ndarray): temporary storage of the current state
         total_step (int): total step numbers
         episode_step (int): step number of the current episode
+        i_episode (int): current episode number
 
     """
 
@@ -72,20 +73,26 @@ class Agent(AbstractAgent):
         self.noise = noise
         self.total_step = 0
         self.episode_step = 0
+        self.i_episode = 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(
-                hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"]
+                self.hyper_params["BUFFER_SIZE"], self.hyper_params["BATCH_SIZE"]
             )
 
     def select_action(self, state: np.ndarray) -> np.ndarray:
         """Select an action from the input space."""
         self.curr_state = state
+        state = self._preprocess_state(state)
 
         # if initial random action should be conducted
         if (
@@ -94,7 +101,6 @@ class Agent(AbstractAgent):
         ):
             return self.env.action_space.sample()
 
-        state = torch.FloatTensor(state).to(device)
         selected_action = self.actor(state)
 
         if not self.args.test:
@@ -103,6 +109,11 @@ class Agent(AbstractAgent):
 
         return selected_action.detach().cpu().numpy()
 
+    def _preprocess_state(self, state: np.ndarray) -> torch.Tensor:
+        """Preprocess state so that actor selects an action."""
+        state = torch.FloatTensor(state).to(device)
+        return state
+
     def step(self, action: np.ndarray) -> Tuple[np.ndarray, np.float64, bool]:
         """Take an action and return the response of the env."""
         self.total_step += 1
@@ -115,10 +126,15 @@ class Agent(AbstractAgent):
             done_bool = (
                 False if self.episode_step == self.args.max_episode_steps else done
             )
-            self.memory.add(self.curr_state, action, reward, next_state, done_bool)
+            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: Tuple[np.ndarray, ...]):
+        """Add 1 step and n step transitions to memory."""
+        self.memory.add(*transition)
+
     def update_model(
         self,
         experiences: Tuple[
@@ -221,7 +237,7 @@ class Agent(AbstractAgent):
             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):
+        for self.i_episode in range(1, self.args.episode_num + 1):
             state = self.env.reset()
             done = False
             score = 0
@@ -229,7 +245,7 @@ class Agent(AbstractAgent):
             loss_episode = list()
 
             while not done:
-                if self.args.render and i_episode >= self.args.render_after:
+                if self.args.render and self.i_episode >= self.args.render_after:
                     self.env.render()
 
                 action = self.select_action(state)
@@ -246,10 +262,10 @@ class Agent(AbstractAgent):
             # logging
             if loss_episode:
                 avg_loss = np.vstack(loss_episode).mean(axis=0)
-                self.write_log(i_episode, avg_loss, score)
+                self.write_log(self.i_episode, avg_loss, score)
 
-            if i_episode % self.args.save_period == 0:
-                self.save_params(i_episode)
+            if self.i_episode % self.args.save_period == 0:
+                self.save_params(self.i_episode)
 
         # termination
         self.env.close()

scripts/algorithms/fd/ddpg_agent.py

@@ -0,0 +1,176 @@
+# -*- coding: utf-8 -*-
+"""DDPGfD agent using demo agent 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
+         https://arxiv.org/pdf/1707.08817.pdf
+"""
+
+import pickle
+from typing import List, Tuple
+
+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.ddpg.agent import Agent as DDPGAgent
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+class Agent(DDPGAgent):
+    """ActorCritic interacting with environment.
+
+    Attributes:
+        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
+            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 for a single step
+            self.beta = self.hyper_params["PER_BETA"]
+            self.memory = PrioritizedReplayBufferfD(
+                self.hyper_params["BUFFER_SIZE"],
+                self.hyper_params["BATCH_SIZE"],
+                demo=list(demos),
+                alpha=self.hyper_params["PER_ALPHA"],
+                epsilon_d=self.hyper_params["PER_EPS_DEMO"],
+            )
+
+    def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]):
+        """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: Tuple[torch.Tensor, ...], gamma: float
+    ) -> torch.Tensor:
+        """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
+        next_actions = self.actor_target(next_states)
+        next_states_actions = torch.cat((next_states, next_actions), dim=-1)
+        next_values = self.critic_target(next_states_actions)
+        curr_returns = rewards + gamma * next_values * masks
+        curr_returns = curr_returns.to(device).detach()
+
+        # train critic
+        values = self.critic(torch.cat((states, actions), dim=-1))
+        critic_loss_element_wise = (values - curr_returns).pow(2)
+
+        return critic_loss_element_wise
+
+    def update_model(
+        self,
+        experiences: Tuple[
+            torch.Tensor,
+            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."""
+        # NOTE This is for old update_model() interface.
+        # experiences_1 = self.memory.sample(self.beta)
+        experiences_1 = experiences
+        states, actions = experiences_1[:2]
+        weights, indices, eps_d = experiences_1[-3:]
+        gamma = self.hyper_params["GAMMA"]
+
+        # train critic
+        critic_loss_element_wise = self._get_critic_loss(experiences_1, gamma)
+        critic_loss = torch.mean(critic_loss_element_wise * weights)
+
+        if self.use_n_step:
+            experiences_n = self.memory_n.sample(indices)
+            gamma = gamma ** self.hyper_params["N_STEP"]
+            critic_loss_n_element_wise = self._get_critic_loss(experiences_n, gamma)
+            # to update loss and priorities
+            lambda1 = self.hyper_params["LAMBDA1"]
+            critic_loss_element_wise += critic_loss_n_element_wise * lambda1
+            critic_loss = torch.mean(critic_loss_element_wise * 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
+        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)
+
+        # increase beta
+        fraction = min(float(self.i_episode) / self.args.episode_num, 1.0)
+        self.beta = self.beta + fraction * (1.0 - self.beta)
+
+        return actor_loss.data, critic_loss.data
+
+    def pretrain(self):
+        """Pretraining steps."""
+        pretrain_loss = list()
+        print("[INFO] Pre-Train %d step." % 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)

scripts/algorithms/fd/sac_agent.py

@@ -0,0 +1,226 @@
+# -*- coding: utf-8 -*-
+"""SAC agent from demonstration for episodic tasks in OpenAI Gym.
+
+- Author: Curt Park
+- Contact: curt.park@medipixel.io
+- Paper: https://arxiv.org/pdf/1801.01290.pdf
+         https://arxiv.org/pdf/1812.05905.pdf
+         https://arxiv.org/pdf/1511.05952.pdf
+         https://arxiv.org/pdf/1707.08817.pdf
+"""
+
+import pickle
+from typing import List, Tuple
+
+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.sac.agent import Agent as SACAgent
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+class Agent(SACAgent):
+    """SAC 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
+            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: Tuple[np.ndarray, ...]):
+        """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)
+
+    # pylint: disable=too-many-statements
+    def update_model(
+        self,
+        experiences: Tuple[
+            torch.Tensor,
+            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, indices, eps_d = (
+            experiences
+        )
+        new_actions, log_prob, pre_tanh_value, mu, std = self.actor(states)
+
+        # train alpha
+        if self.hyper_params["AUTO_ENTROPY_TUNING"]:
+            alpha_loss = torch.mean(
+                (-self.log_alpha * (log_prob + self.target_entropy).detach()) * weights
+            )
+
+            self.alpha_optimizer.zero_grad()
+            alpha_loss.backward()
+            self.alpha_optimizer.step()
+
+            alpha = self.log_alpha.exp()
+        else:
+            alpha_loss = torch.zeros(1)
+            alpha = self.hyper_params["W_ENTROPY"]
+
+        # Q function loss
+        masks = 1 - dones
+        gamma = self.hyper_params["GAMMA"]
+        q_1_pred = self.qf_1(states, actions)
+        q_2_pred = self.qf_2(states, actions)
+        v_target = self.vf_target(next_states)
+        q_target = rewards + self.hyper_params["GAMMA"] * v_target * masks
+        qf_1_loss = torch.mean((q_1_pred - q_target.detach()).pow(2) * weights)
+        qf_2_loss = torch.mean((q_2_pred - q_target.detach()).pow(2) * weights)
+
+        if self.use_n_step:
+            experiences_n = self.memory_n.sample(indices)
+            _, _, rewards, next_states, dones = experiences_n
+            gamma = gamma ** self.hyper_params["N_STEP"]
+            lambda1 = self.hyper_params["LAMBDA1"]
+            masks = 1 - dones
+
+            v_target = self.vf_target(next_states)
+            q_target = rewards + gamma * v_target * masks
+            qf_1_loss_n = torch.mean((q_1_pred - q_target.detach()).pow(2) * weights)
+            qf_2_loss_n = torch.mean((q_2_pred - q_target.detach()).pow(2) * weights)
+
+            # to update loss and priorities
+            qf_1_loss = qf_1_loss + qf_1_loss_n * lambda1
+            qf_2_loss = qf_2_loss + qf_2_loss_n * lambda1
+
+        # V function loss
+        v_pred = self.vf(states)
+        q_pred = torch.min(
+            self.qf_1(states, new_actions), self.qf_2(states, new_actions)
+        )
+        v_target = (q_pred - alpha * log_prob).detach()
+        vf_loss_element_wise = (v_pred - v_target).pow(2)
+        vf_loss = torch.mean(vf_loss_element_wise * weights)
+
+        # train Q functions
+        self.qf_1_optimizer.zero_grad()
+        qf_1_loss.backward()
+        self.qf_1_optimizer.step()
+
+        self.qf_2_optimizer.zero_grad()
+        qf_2_loss.backward()
+        self.qf_2_optimizer.step()
+
+        # train V function
+        self.vf_optimizer.zero_grad()
+        vf_loss.backward()
+        self.vf_optimizer.step()
+
+        if self.total_step % self.hyper_params["DELAYED_UPDATE"] == 0:
+            # actor loss
+            advantage = q_pred - v_pred.detach()
+            actor_loss_element_wise = alpha * log_prob - advantage
+            actor_loss = torch.mean(actor_loss_element_wise * weights)
+
+            # regularization
+            mean_reg = self.hyper_params["W_MEAN_REG"] * mu.pow(2).mean()
+            std_reg = self.hyper_params["W_STD_REG"] * std.pow(2).mean()
+            pre_activation_reg = self.hyper_params["W_PRE_ACTIVATION_REG"] * (
+                pre_tanh_value.pow(2).sum(dim=-1).mean()
+            )
+            actor_reg = mean_reg + std_reg + pre_activation_reg
+
+            # actor loss + regularization
+            actor_loss += actor_reg
+
+            # train actor
+            self.actor_optimizer.zero_grad()
+            actor_loss.backward()
+            self.actor_optimizer.step()
+
+            # update target networks
+            common_utils.soft_update(self.vf, self.vf_target, self.hyper_params["TAU"])
+
+            # update priorities
+            new_priorities = vf_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)
+
+            # increase beta
+            fraction = min(float(self.i_episode) / self.args.episode_num, 1.0)
+            self.beta = self.beta + fraction * (1.0 - self.beta)
+        else:
+            actor_loss = torch.zeros(1)
+
+        return (
+            actor_loss.data,
+            qf_1_loss.data,
+            qf_2_loss.data,
+            vf_loss.data,
+            alpha_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"]
+                )

scripts/algorithms/fd/td3_agent.py

@@ -0,0 +1,209 @@
+# -*- 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
+from typing import List, Tuple
+
+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
+            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: Tuple[np.ndarray, ...]):
+        """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: Tuple[torch.Tensor, ...], gamma: float
+    ) -> torch.Tensor:
+        """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 + (self.hyper_params["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: Tuple[
+            torch.Tensor,
+            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, 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"]
+                )

scripts/algorithms/sac/agent.py

@@ -46,6 +46,7 @@ class Agent(AbstractAgent):
         hyper_params (dict): hyper-parameters
         total_step (int): total step numbers
         episode_step (int): step number of the current episode
+        i_episode (int): current episode number
 
     """
 
@@ -78,6 +79,7 @@ class Agent(AbstractAgent):
         self.curr_state = np.zeros((1,))
         self.total_step = 0
         self.episode_step = 0
+        self.i_episode = 0
 
         # automatic entropy tuning
         if self.hyper_params["AUTO_ENTROPY_TUNING"]:
@@ -91,15 +93,20 @@ class Agent(AbstractAgent):
         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(
-                hyper_params["BUFFER_SIZE"], hyper_params["BATCH_SIZE"]
+                self.hyper_params["BUFFER_SIZE"], self.hyper_params["BATCH_SIZE"]
             )
 
     def select_action(self, state: np.ndarray) -> np.ndarray:
         """Select an action from the input space."""
         self.curr_state = state
+        state = self._preprocess_state(state)
 
         # if initial random action should be conducted
         if (
@@ -108,7 +115,6 @@ class Agent(AbstractAgent):
         ):
             return self.env.action_space.sample()
 
-        state = torch.FloatTensor(state).to(device)
         if self.args.test:
             _, _, _, selected_action, _ = self.actor(state)
         else:
@@ -116,6 +122,11 @@ class Agent(AbstractAgent):
 
         return selected_action.detach().cpu().numpy()
 
+    def _preprocess_state(self, state: np.ndarray) -> torch.Tensor:
+        """Preprocess state so that actor selects an action."""
+        state = torch.FloatTensor(state).to(device)
+        return state
+
     def step(self, action: np.ndarray) -> Tuple[np.ndarray, np.float64, bool]:
         """Take an action and return the response of the env."""
         self.total_step += 1
@@ -128,10 +139,15 @@ class Agent(AbstractAgent):
             done_bool = (
                 False if self.episode_step == self.args.max_episode_steps else done
             )
-            self.memory.add(self.curr_state, action, reward, next_state, done_bool)
+            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: Tuple[np.ndarray, ...]):
+        """Add 1 step and n step transitions to memory."""
+        self.memory.add(*transition)
+
     def update_model(
         self,
         experiences: Tuple[
@@ -308,7 +324,7 @@ class Agent(AbstractAgent):
             wandb.config.update(self.hyper_params)
             wandb.watch([self.actor, self.vf, self.qf_1, self.qf_2], log="parameters")
 
-        for i_episode in range(1, self.args.episode_num + 1):
+        for self.i_episode in range(1, self.args.episode_num + 1):
             state = self.env.reset()
             done = False
             score = 0
@@ -316,7 +332,7 @@ class Agent(AbstractAgent):
             loss_episode = list()
 
             while not done:
-                if self.args.render and i_episode >= self.args.render_after:
+                if self.args.render and self.i_episode >= self.args.render_after:
                     self.env.render()
 
                 action = self.select_action(state)
@@ -335,11 +351,11 @@ class Agent(AbstractAgent):
             if loss_episode:
                 avg_loss = np.vstack(loss_episode).mean(axis=0)
                 self.write_log(
-                    i_episode, avg_loss, score, self.hyper_params["DELAYED_UPDATE"]
+                    self.i_episode, avg_loss, score, self.hyper_params["DELAYED_UPDATE"]
                 )
 
-            if i_episode % self.args.save_period == 0:
-                self.save_params(i_episode)
+            if self.i_episode % self.args.save_period == 0:
+                self.save_params(self.i_episode)
 
         # termination
         self.env.close()

scripts/algorithms/td3/agent.py

@@ -65,8 +65,9 @@ class Agent(AbstractAgent):
 
         """
         AbstractAgent.__init__(self, env, args)
-        self.actor, self.actor_target, self.critic1, self.critic1_target, \
-            self.critic2, self.critic2_target = models
+        self.actor, self.actor_target, self.critic1, self.critic1_target, self.critic2, self.critic2_target = (  # noqa: B950
+            models
+        )
         self.actor_optim, self.critic_optim = optims
         self.hyper_params = hyper_params
         self.exploration_noise, self.target_policy_noise = noises
@@ -78,10 +79,15 @@ class Agent(AbstractAgent):
         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._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: np.ndarray) -> np.ndarray:
         """Select an action from the input space."""
@@ -104,12 +110,25 @@ class Agent(AbstractAgent):
 
     def step(self, action: torch.Tensor) -> Tuple[np.ndarray, np.float64, bool]:
         """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)
 
-        self.memory.add(self.curr_state, action, reward, next_state, done)
+        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: Tuple[np.ndarray, ...]):
+        """Add 1 step and n step transitions to memory."""
+        self.memory.add(*transition)
+
     def update_model(
         self,
         experiences: Tuple[
@@ -137,7 +156,9 @@ class Agent(AbstractAgent):
             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()
+        target_values = (
+            rewards + (self.hyper_params["GAMMA"] * target_values * masks).detach()
+        )
 
         # train critic
         values1 = self.critic1(torch.cat((states, actions), dim=-1))
@@ -206,8 +227,8 @@ class Agent(AbstractAgent):
 
         print(
             "[INFO] total_steps: %d episode: %d total score: %d, total loss: %f\n"
-            "actor_loss: %.3f critic_loss: %.3f\n"
-            % (self.total_steps, i, score, total_loss, loss[0], loss[1])
+            "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:
@@ -243,8 +264,6 @@ class Agent(AbstractAgent):
 
                 action = self.select_action(state)
                 next_state, reward, done = self.step(action)
-                self.total_steps += 1
-                self.episode_steps += 1
 
                 if len(self.memory) >= self.hyper_params["BATCH_SIZE"]:
                     experiences = self.memory.sample()

scripts/examples/lunarlander_continuous_v2/ddpgfd.py

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+"""Run module for DDPGfD on LunarLanderContinuous-v2.
+
+- Author: Curt Park
+- Contact: curt.park@medipixel.io
+"""
+
+import argparse
+
+import gym
+import torch
+import torch.optim as optim
+
+from algorithms.common.networks.mlp import MLP
+from algorithms.common.noise import OUNoise
+from algorithms.fd.ddpg_agent import Agent
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# hyper parameters
+hyper_params = {
+    "N_STEP": 1,
+    "GAMMA": 0.99,
+    "TAU": 5e-3,
+    "BUFFER_SIZE": int(1e5),
+    "BATCH_SIZE": 64,
+    "LR_ACTOR": 3e-4,
+    "LR_CRITIC": 3e-4,
+    "OU_NOISE_THETA": 0.0,
+    "OU_NOISE_SIGMA": 0.0,
+    "PRETRAIN_STEP": 0,
+    "MULTIPLE_LEARN": 2,  # multiple learning updates
+    "LAMBDA1": 1.0,  # N-step return weight
+    "LAMBDA2": 1e-5,  # l2 regularization weight
+    "LAMBDA3": 1.0,  # actor loss contribution of prior weight
+    "PER_ALPHA": 0.3,
+    "PER_BETA": 1.0,
+    "PER_EPS": 1e-6,
+    "PER_EPS_DEMO": 1.0,
+    "INITIAL_RANDOM_ACTION": int(5e3),
+}
+
+
+def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int):
+    """Run training or test.
+
+    Args:
+        env (gym.Env): openAI Gym environment with continuous action space
+        args (argparse.Namespace): arguments including training settings
+        state_dim (int): dimension of states
+        action_dim (int): dimension of actions
+
+    """
+    hidden_sizes_actor = [256, 256]
+    hidden_sizes_critic = [256, 256]
+
+    # create actor
+    actor = MLP(
+        input_size=state_dim,
+        output_size=action_dim,
+        hidden_sizes=hidden_sizes_actor,
+        output_activation=torch.tanh,
+    ).to(device)
+
+    actor_target = MLP(
+        input_size=state_dim,
+        output_size=action_dim,
+        hidden_sizes=hidden_sizes_actor,
+        output_activation=torch.tanh,
+    ).to(device)
+    actor_target.load_state_dict(actor.state_dict())
+
+    # create critic
+    critic = MLP(
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
+    ).to(device)
+
+    critic_target = MLP(
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
+    ).to(device)
+    critic_target.load_state_dict(critic.state_dict())
+
+    # create optimizer
+    actor_optim = optim.Adam(
+        actor.parameters(),
+        lr=hyper_params["LR_ACTOR"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+
+    critic_optim = optim.Adam(
+        critic.parameters(),
+        lr=hyper_params["LR_CRITIC"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+
+    # noise
+    noise = OUNoise(
+        action_dim,
+        theta=hyper_params["OU_NOISE_THETA"],
+        sigma=hyper_params["OU_NOISE_SIGMA"],
+    )
+
+    # make tuples to create an agent
+    models = (actor, actor_target, critic, critic_target)
+    optims = (actor_optim, critic_optim)
+
+    # create an agent
+    agent = Agent(env, args, hyper_params, models, optims, noise)
+
+    # run
+    if args.test:
+        agent.test()
+    else:
+        agent.train()

scripts/examples/lunarlander_continuous_v2/sacfd.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+"""Run module for SACfD on LunarLanderContinuous-v2.
+
+- Author: Curt Park
+- Contact: curt.park@medipixel.io
+"""
+
+import argparse
+
+import gym
+import numpy as np
+import torch
+import torch.optim as optim
+
+from algorithms.common.networks.mlp import MLP, FlattenMLP, TanhGaussianDistParams
+from algorithms.fd.sac_agent import Agent
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# hyper parameters
+hyper_params = {
+    "N_STEP": 3,
+    "GAMMA": 0.99,
+    "TAU": 1e-3,
+    "BUFFER_SIZE": int(1e5),
+    "BATCH_SIZE": 64,
+    "AUTO_ENTROPY_TUNING": True,
+    "LR_ACTOR": 3e-4,
+    "LR_VF": 3e-4,
+    "LR_QF1": 3e-4,
+    "LR_QF2": 3e-4,
+    "LR_ENTROPY": 3e-4,
+    "W_ENTROPY": 1e-3,
+    "W_MEAN_REG": 1e-3,
+    "W_STD_REG": 1e-3,
+    "W_PRE_ACTIVATION_REG": 0.0,
+    "DELAYED_UPDATE": 2,
+    "PRETRAIN_STEP": 100,
+    "MULTIPLE_LEARN": 2,  # multiple learning updates
+    "LAMBDA1": 1.0,  # N-step return weight
+    "LAMBDA2": 1e-5,  # l2 regularization weight
+    "LAMBDA3": 1.0,  # actor loss contribution of prior weight
+    "PER_ALPHA": 0.6,
+    "PER_BETA": 0.4,
+    "PER_EPS": 1e-6,
+    "PER_EPS_DEMO": 1.0,
+    "INITIAL_RANDOM_ACTION": int(5e3),
+}
+
+
+def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int):
+    """Run training or test.
+
+    Args:
+        env (gym.Env): openAI Gym environment with continuous action space
+        args (argparse.Namespace): arguments including training settings
+        state_dim (int): dimension of states
+        action_dim (int): dimension of actions
+
+    """
+    hidden_sizes_actor = [256, 256]
+    hidden_sizes_vf = [256, 256]
+    hidden_sizes_qf = [256, 256]
+
+    # target entropy
+    target_entropy = -np.prod((action_dim,)).item()  # heuristic
+
+    # create actor
+    actor = TanhGaussianDistParams(
+        input_size=state_dim, output_size=action_dim, hidden_sizes=hidden_sizes_actor
+    ).to(device)
+
+    # create v_critic
+    vf = MLP(input_size=state_dim, output_size=1, hidden_sizes=hidden_sizes_vf).to(
+        device
+    )
+    vf_target = MLP(
+        input_size=state_dim, output_size=1, hidden_sizes=hidden_sizes_vf
+    ).to(device)
+    vf_target.load_state_dict(vf.state_dict())
+
+    # create q_critic
+    qf_1 = FlattenMLP(
+        input_size=state_dim + action_dim, output_size=1, hidden_sizes=hidden_sizes_qf
+    ).to(device)
+    qf_2 = FlattenMLP(
+        input_size=state_dim + action_dim, output_size=1, hidden_sizes=hidden_sizes_qf
+    ).to(device)
+
+    # create optimizers
+    actor_optim = optim.Adam(
+        actor.parameters(),
+        lr=hyper_params["LR_ACTOR"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+    vf_optim = optim.Adam(
+        vf.parameters(), lr=hyper_params["LR_VF"], weight_decay=hyper_params["LAMBDA2"]
+    )
+    qf_1_optim = optim.Adam(
+        qf_1.parameters(),
+        lr=hyper_params["LR_QF1"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+    qf_2_optim = optim.Adam(
+        qf_2.parameters(),
+        lr=hyper_params["LR_QF2"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+
+    # make tuples to create an agent
+    models = (actor, vf, vf_target, qf_1, qf_2)
+    optims = (actor_optim, vf_optim, qf_1_optim, qf_2_optim)
+
+    # create an agent
+    agent = Agent(env, args, hyper_params, models, optims, target_entropy)
+
+    # run
+    if args.test:
+        agent.test()
+    else:
+        agent.train()

scripts/examples/lunarlander_continuous_v2/td3.py

@@ -65,25 +65,29 @@ def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int)
 
     # create critic1
     critic1 = MLP(
-        input_size=state_dim + action_dim, output_size=1,
-        hidden_sizes=hidden_sizes_critic
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
     ).to(device)
 
     critic1_target = MLP(
-        input_size=state_dim + action_dim, output_size=1,
-        hidden_sizes=hidden_sizes_critic
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
     ).to(device)
     critic1_target.load_state_dict(critic1.state_dict())
 
     # create critic2
     critic2 = MLP(
-        input_size=state_dim + action_dim, output_size=1,
-        hidden_sizes=hidden_sizes_critic
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
     ).to(device)
 
     critic2_target = MLP(
-        input_size=state_dim + action_dim, output_size=1,
-        hidden_sizes=hidden_sizes_critic
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
     ).to(device)
     critic2_target.load_state_dict(critic2.state_dict())
 

scripts/examples/lunarlander_continuous_v2/td3fd.py

@@ -0,0 +1,147 @@
+# -*- coding: utf-8 -*-
+"""Run module for SACfD on LunarLanderContinuous-v2.
+
+- Author: Seungjae Ryan Lee
+- Contact: seungjaeryanlee@gmail.com
+"""
+
+import argparse
+
+import gym
+import torch
+import torch.optim as optim
+
+from algorithms.common.networks.mlp import MLP
+from algorithms.common.noise import GaussianNoise
+from algorithms.fd.td3_agent import Agent
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# hyper parameters
+# TODO Tune hyperparameters on LunarLander-v2
+
+# hyper parameters
+hyper_params = {
+    "N_STEP": 3,
+    "GAMMA": 0.99,
+    "TAU": 1e-3,
+    "BUFFER_SIZE": int(1e5),
+    "BATCH_SIZE": 64,
+    "LR_ACTOR": 3e-4,
+    "LR_CRITIC": 3e-4,
+    "EXPLORATION_NOISE": 0.1,
+    "TARGET_POLICY_NOISE": 0.2,
+    "TARGET_POLICY_NOISE_CLIP": 0.5,
+    "POLICY_UPDATE_FREQ": 2,
+    "INITIAL_RANDOM_ACTIONS": 5e3,
+    "PRETRAIN_STEP": 100,
+    "MULTIPLE_LEARN": 2,  # multiple learning updates
+    "LAMBDA1": 1.0,  # N-step return weight
+    "LAMBDA2": 1e-5,  # l2 regularization weight
+    "LAMBDA3": 1.0,  # actor loss contribution of prior weight
+    "PER_ALPHA": 0.3,
+    "PER_BETA": 1.0,
+    "PER_EPS": 1e-6,
+    "PER_EPS_DEMO": 1.0,
+}
+
+
+def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int):
+    """Run training or test.
+
+    Args:
+        env (gym.Env): openAI Gym environment with continuous action space
+        args (argparse.Namespace): arguments including training settings
+        state_dim (int): dimension of states
+        action_dim (int): dimension of actions
+
+    """
+    hidden_sizes_actor = [400, 300]
+    hidden_sizes_critic = [400, 300]
+
+    # create actor
+    actor = MLP(
+        input_size=state_dim,
+        output_size=action_dim,
+        hidden_sizes=hidden_sizes_actor,
+        output_activation=torch.tanh,
+    ).to(device)
+
+    actor_target = MLP(
+        input_size=state_dim,
+        output_size=action_dim,
+        hidden_sizes=hidden_sizes_actor,
+        output_activation=torch.tanh,
+    ).to(device)
+    actor_target.load_state_dict(actor.state_dict())
+
+    # create critic1
+    critic1 = MLP(
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
+    ).to(device)
+
+    critic1_target = MLP(
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
+    ).to(device)
+    critic1_target.load_state_dict(critic1.state_dict())
+
+    # create critic2
+    critic2 = MLP(
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
+    ).to(device)
+
+    critic2_target = MLP(
+        input_size=state_dim + action_dim,
+        output_size=1,
+        hidden_sizes=hidden_sizes_critic,
+    ).to(device)
+    critic2_target.load_state_dict(critic2.state_dict())
+
+    # concat critic parameters to use one optim
+    critic_parameters = list(critic1.parameters()) + list(critic2.parameters())
+
+    # create optimizer
+    actor_optim = optim.Adam(
+        actor.parameters(),
+        lr=hyper_params["LR_ACTOR"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+
+    critic_optim = optim.Adam(
+        critic_parameters,
+        lr=hyper_params["LR_CRITIC"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+
+    # noise
+    exploration_noise = GaussianNoise(
+        action_dim,
+        min_sigma=hyper_params["EXPLORATION_NOISE"],
+        max_sigma=hyper_params["EXPLORATION_NOISE"],
+    )
+
+    target_policy_noise = GaussianNoise(
+        action_dim,
+        min_sigma=hyper_params["TARGET_POLICY_NOISE"],
+        max_sigma=hyper_params["TARGET_POLICY_NOISE"],
+    )
+
+    # make tuples to create an agent
+    models = (actor, actor_target, critic1, critic1_target, critic2, critic2_target)
+    optims = (actor_optim, critic_optim)
+    noises = (exploration_noise, target_policy_noise)
+
+    # create an agent
+    agent = Agent(env, args, hyper_params, models, optims, noises)
+
+    # run
+    if args.test:
+        agent.test()
+    else:
+        agent.train()

scripts/run_lunarlander_continuous.py

@@ -40,6 +40,12 @@ parser.add_argument(
 parser.add_argument("--save-period", type=int, default=100, help="save model period")
 parser.add_argument("--log", action="store_true", help="turn on logging")
 parser.add_argument("--test", action="store_true", help="test mode (no training)")
+parser.add_argument(
+    "--demo-path",
+    type=str,
+    default="data/lunarlander_continuous_demo.pkl",
+    help="demonstration path",
+)
 parser.set_defaults(render=True)
 
 args = parser.parse_args()

scripts/test_lstm.py

@@ -1,13 +1,12 @@
 # -*- coding: utf-8 -*-
 """Simple test case(sine+noise->cos) for LSTM network."""
 
-import numpy as np
 import matplotlib.pyplot as plt
-
+import numpy as np
 import torch
-from torch.utils.data import Dataset, DataLoader
-import torch.optim as optim
 import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
 
 from algorithms.common.networks.lstm import LSTM