Merge branch 'master' into feat/demo_refactoring

scripts/algorithms/common/abstract/her.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+"""Abstract class used for Hindsight Experience Replay.
+- Author: Kh Kim
+- Contact: kh.kim@medipixel.io
+- Paper: https://arxiv.org/pdf/1707.01495.pdf
+"""
+
+from abc import ABCMeta, abstractmethod
+
+import numpy as np
+
+
+class HER(object):
+    """Abstract class for HER (final strategy).
+    Attributes:
+        reward_func (Callable): returns reward from state, action, next_state
+    """
+
+    __metaclass__ = ABCMeta
+
+    def __init__(self, reward_func):
+        """Initialization.
+
+        Args:
+            reward_func (Callable): returns reward from state, action, next_state
+        """
+        self.reward_func = reward_func()
+
+    @abstractmethod
+    def fetch_desired_states_from_demo(self, demo):
+        pass
+
+    @abstractmethod
+    def get_desired_state(self, *args):
+        pass
+
+    @abstractmethod
+    def generate_demo_transitions(self, demo):
+        pass
+
+    @abstractmethod
+    def _get_final_state(self, transition):
+        pass
+
+    def _append_origin_transitions(self, origin_transitions, transition, desired_state):
+        """Append original transitions adding goal state for training."""
+        origin_transitions.append(self._get_transition(transition, desired_state))
+
+    def _append_new_transitions(self, new_transitions, transition, final_state):
+        """Append new transitions made by HER strategy (final) for training."""
+        new_transitions.append(self._get_transition(transition, final_state))
+
+    def _get_transition(self, transition, goal_state):
+        """Get a single transition concatenated with a goal state."""
+        state, action, _, next_state, done = transition
+
+        done = np.array_equal(next_state, goal_state)
+        reward = self.reward_func(transition, goal_state)
+        state = np.concatenate((state, goal_state), axis=-1)
+        next_state = np.concatenate((next_state, goal_state), axis=-1)
+
+        return state, action, reward, next_state, done
+
+    def generate_transitions(
+        self, transitions, desired_state, success_score, is_demo=False
+    ):
+        """Generate new transitions concatenated with desired states."""
+        origin_transitions = list()
+        new_transitions = list()
+        final_state = self._get_final_state(transitions[-1])
+        score = np.sum(np.array(transitions), axis=0)[2]
+
+        for transition in transitions:
+            # process transitions with the initial goal state
+            self._append_origin_transitions(
+                origin_transitions, transition, desired_state
+            )
+
+            # do not need to append new transitions if sum of reward is big enough
+            if not is_demo and score <= success_score:
+                self._append_new_transitions(new_transitions, transition, final_state)
+
+        return origin_transitions + new_transitions

scripts/algorithms/common/abstract/reward_fn.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+"""Abstract class for computing reward.
+- Author: Kh Kim
+- Contact: kh.kim@medipixel.io
+"""
+
+from abc import ABCMeta, abstractmethod
+
+
+class RewardFn(object):
+    """Abstract class for computing reward.
+       New compute_reward class should redefine __call__()
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def __call__(self, transition, goal_state):
+        pass

scripts/algorithms/sac/agent.py

@@ -8,6 +8,7 @@
 """
 
 import os
+import pickle
 
 import numpy as np
 import torch
@@ -44,10 +45,11 @@ class Agent(AbstractAgent):
         total_step (int): total step numbers
         episode_step (int): step number of the current episode
         i_episode (int): current episode number
+        her (HER): hinsight experience replay
 
     """
 
-    def __init__(self, env, args, hyper_params, models, optims, target_entropy):
+    def __init__(self, env, args, hyper_params, models, optims, target_entropy, her):
         """Initialization.
 
         Args:
@@ -57,6 +59,7 @@ class Agent(AbstractAgent):
             models (tuple): models including actor and critic
             optims (tuple): optimizers for actor and critic
             target_entropy (float): target entropy for the inequality constraint
+            her (HER): hinsight experience replay
 
         """
         AbstractAgent.__init__(self, env, args)
@@ -69,6 +72,7 @@ class Agent(AbstractAgent):
         self.total_step = 0
         self.episode_step = 0
         self.i_episode = 0
+        self.her = her
 
         # automatic entropy tuning
         if self.hyper_params["AUTO_ENTROPY_TUNING"]:
@@ -92,6 +96,50 @@ class Agent(AbstractAgent):
                 self.hyper_params["BUFFER_SIZE"], self.hyper_params["BATCH_SIZE"]
             )
 
+        # HER
+        if self.hyper_params["USE_HER"]:
+            # load demo replay memory
+            with open(self.args.demo_path, "rb") as f:
+                demo = pickle.load(f)
+
+            if self.hyper_params["DESIRED_STATES_FROM_DEMO"]:
+                self.her.fetch_desired_states_from_demo(demo)
+
+            self.transitions_epi = list()
+            self.desired_state = np.zeros((1,))
+            demo = self.her.generate_demo_transitions(demo)
+
+        if not self.args.test:
+            # Replay buffers
+            self.memory = ReplayBuffer(
+                self.hyper_params["BUFFER_SIZE"], self.hyper_params["BATCH_SIZE"]
+            )
+
+    def _preprocess_state(self, state):
+        """Preprocess state so that actor selects an action."""
+        if self.hyper_params["USE_HER"]:
+            self.desired_state = self.her.get_desired_state()
+            state = np.concatenate((state, self.desired_state), axis=-1)
+        state = torch.FloatTensor(state).to(device)
+        return state
+
+    def _add_transition_to_memory(self, transition):
+        """Add 1 step and n step transitions to memory."""
+        if self.hyper_params["USE_HER"]:
+            self.transitions_epi.append(transition)
+            done = transition[-1] or self.episode_step == self.args.max_episode_steps
+            if done:
+                # insert generated transitions if the episode is done
+                transitions = self.her.generate_transitions(
+                    self.transitions_epi,
+                    self.desired_state,
+                    self.hyper_params["SUCCESS_SCORE"],
+                )
+                self.memory.extend(transitions)
+                self.transitions_epi = list()
+        else:
+            self.memory.add(*transition)
+
     def select_action(self, state):
         """Select an action from the input space."""
         self.curr_state = state
@@ -111,11 +159,6 @@ class Agent(AbstractAgent):
 
         return selected_action.detach().cpu().numpy()
 
-    def _preprocess_state(self, state):
-        """Preprocess state so that actor selects an action."""
-        state = torch.FloatTensor(state).to(device)
-        return state
-
     def step(self, action):
         """Take an action and return the response of the env."""
         self.total_step += 1
@@ -133,10 +176,6 @@ class Agent(AbstractAgent):
 
         return next_state, reward, done
 
-    def _add_transition_to_memory(self, transition):
-        """Add 1 step and n step transitions to memory."""
-        self.memory.add(*transition)
-
     def update_model(self, experiences):
         """Train the model after each episode."""
         states, actions, rewards, next_states, dones = experiences
@@ -304,6 +343,7 @@ class Agent(AbstractAgent):
         if self.args.log:
             wandb.init()
             wandb.config.update(self.hyper_params)
+            wandb.config.update(vars(self.args))
             wandb.watch([self.actor, self.vf, self.qf_1, self.qf_2], log="parameters")
 
         for self.i_episode in range(1, self.args.episode_num + 1):

scripts/algorithms/td3/agent.py

@@ -234,6 +234,7 @@ class Agent(AbstractAgent):
         if self.args.log:
             wandb.init()
             wandb.config.update(self.hyper_params)
+            wandb.config.update(vars(self.args))
             wandb.watch([self.actor, self.critic1, self.critic2], log="parameters")
 
         for i_episode in range(1, self.args.episode_num + 1):