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
2026-08-29 11:20:40 +08:00 · 2019-03-14 11:06:54 +09:00
parent ee014e5a93
commit ca5c99bc41
17 changed files with 1377 additions and 49 deletions
@@ -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:
@@ -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)
@@ -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_
@@ -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
@@ -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()
@@ -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)
@@ -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"]
+                )
@@ -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"]
+                )
@@ -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()
@@ -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()