Spaces:
Configuration error
Configuration error
File size: 3,052 Bytes
b89a51c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import numpy as np
class ExperienceReplay:
def __init__(self,
num_frame_stack=4,
capacity=int(1e5),
pic_size=(96, 96)
):
self.num_frame_stack = num_frame_stack
self.capacity = capacity
self.pic_size = pic_size
self.counter = 0
self.frame_window = None
self.init_caches()
self.expecting_new_episode = True
def add_experience(self, frame, action, done, reward):
assert self.frame_window is not None, "start episode first"
self.counter += 1
frame_idx = self.counter % self.max_frame_cache
exp_idx = (self.counter - 1) % self.capacity
self.prev_states[exp_idx] = self.frame_window
self.frame_window = np.append(self.frame_window[1:], frame_idx)
self.next_states[exp_idx] = self.frame_window
self.actions[exp_idx] = action
self.is_done[exp_idx] = done
self.frames[frame_idx] = frame
self.rewards[exp_idx] = reward
if done:
self.expecting_new_episode = True
def start_new_episode(self, frame):
# it should be okay not to increment counter here
# because episode ending frames are not used
assert self.expecting_new_episode, "previous episode didn't end yet"
frame_idx = self.counter % self.max_frame_cache
self.frame_window = np.repeat(frame_idx, self.num_frame_stack)
self.frames[frame_idx] = frame
self.expecting_new_episode = False
def sample_mini_batch(self, n):
count = min(self.capacity, self.counter)
batchidx = np.random.randint(count, size=n)
prev_frames = self.frames[self.prev_states[batchidx]]
next_frames = self.frames[self.next_states[batchidx]]
prev_frames = np.moveaxis(prev_frames, 1, -1)
next_frames = np.moveaxis(next_frames, 1, -1)
return {
"reward": self.rewards[batchidx],
"prev_state": prev_frames,
"next_state": next_frames,
"actions": self.actions[batchidx],
"done_mask": self.is_done[batchidx]
}
def current_state(self):
# assert not self.expecting_new_episode, "start new episode first"'
assert self.frame_window is not None, "do something first"
sf = self.frames[self.frame_window]
sf = np.moveaxis(sf, 0, -1)
return sf
def init_caches(self):
self.rewards = np.zeros(self.capacity, dtype="float32")
self.prev_states = -np.ones((self.capacity, self.num_frame_stack),
dtype="int32")
self.next_states = -np.ones((self.capacity, self.num_frame_stack),
dtype="int32")
self.is_done = -np.ones(self.capacity, "int32")
self.actions = -np.ones(self.capacity, dtype="int32")
self.max_frame_cache = self.capacity + 2 * self.num_frame_stack + 1
self.frames = -np.ones((self.max_frame_cache,) + self.pic_size, dtype="float32")
|