Implementation of hybrid reward architecture

component/replay.py

@@ -49,6 +49,50 @@ class Replay:
                 self.next_states[sampled_indices],
                 self.terminals[sampled_indices]]
 
+class HybridRewardReplay:
+    def __init__(self, memory_size, batch_size, dtype=np.float32):
+        self.memory_size = memory_size
+        self.batch_size = batch_size
+        self.dtype = dtype
+
+        self.states = None
+        self.actions = np.empty(self.memory_size, dtype=np.int8)
+        self.rewards = None
+        self.next_states = None
+        self.terminals = np.empty(self.memory_size, dtype=np.int8)
+
+        self.pos = 0
+        self.full = False
+
+
+    def feed(self, experience):
+        state, action, reward, next_state, done = experience
+
+        if self.states is None:
+            self.rewards = np.empty((self.memory_size, ) + reward.shape, dtype=self.dtype)
+            self.states = np.empty((self.memory_size, ) + state.shape, dtype=self.dtype)
+            self.next_states = np.empty((self.memory_size, ) + state.shape, dtype=self.dtype)
+
+        self.states[self.pos][:] = state
+        self.actions[self.pos] = action
+        self.rewards[self.pos][:] = reward
+        self.next_states[self.pos][:] = next_state
+        self.terminals[self.pos] = done
+
+        self.pos += 1
+        if self.pos == self.memory_size:
+            self.full = True
+            self.pos = 0
+
+    def sample(self):
+        upper_bound = self.memory_size if self.full else self.pos
+        sampled_indices = np.random.randint(0, upper_bound, size=self.batch_size)
+        return [self.states[sampled_indices],
+                self.actions[sampled_indices],
+                self.rewards[sampled_indices],
+                self.next_states[sampled_indices],
+                self.terminals[sampled_indices]]
+
 class HighDimActionReplay:
     def __init__(self, memory_size, batch_size, dtype=np.float32):
         self.memory_size = memory_size
@@ -91,4 +135,5 @@ class HighDimActionReplay:
                 self.actions[sampled_indices],
                 self.rewards[sampled_indices],
                 self.next_states[sampled_indices],
-                self.terminals[sampled_indices]]
+                self.terminals[sampled_indices]]
+