93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
import numpy as np
|
|
from tianshou.data.batch import Batch
|
|
|
|
|
|
class ReplayBuffer(object):
|
|
"""docstring for ReplayBuffer"""
|
|
|
|
def __init__(self, size):
|
|
super().__init__()
|
|
self._maxsize = size
|
|
self.reset()
|
|
|
|
def __len__(self):
|
|
return self._size
|
|
|
|
def _add_to_buffer(self, name, inst):
|
|
if inst is None:
|
|
return
|
|
if self.__dict__.get(name, None) is None:
|
|
if isinstance(inst, np.ndarray):
|
|
self.__dict__[name] = np.zeros([self._maxsize, *inst.shape])
|
|
elif isinstance(inst, dict):
|
|
self.__dict__[name] = np.array(
|
|
[{} for _ in range(self._maxsize)])
|
|
else: # assume `inst` is a number
|
|
self.__dict__[name] = np.zeros([self._maxsize])
|
|
self.__dict__[name][self._index] = inst
|
|
|
|
def update(self, buffer):
|
|
for i in range(len(buffer)):
|
|
self.add(
|
|
buffer.obs[i], buffer.act[i], buffer.rew[i],
|
|
buffer.done[i], buffer.obs_next[i], buffer.info[i])
|
|
|
|
def add(self, obs, act, rew, done, obs_next=0, info={}, weight=None):
|
|
'''
|
|
weight: importance weights, disabled here
|
|
'''
|
|
assert isinstance(info, dict),\
|
|
'You should return a dict in the last argument of env.step().'
|
|
self._add_to_buffer('obs', obs)
|
|
self._add_to_buffer('act', act)
|
|
self._add_to_buffer('rew', rew)
|
|
self._add_to_buffer('done', done)
|
|
self._add_to_buffer('obs_next', obs_next)
|
|
self._add_to_buffer('info', info)
|
|
self._size = min(self._size + 1, self._maxsize)
|
|
self._index = (self._index + 1) % self._maxsize
|
|
|
|
def reset(self):
|
|
self._index = self._size = 0
|
|
self.indice = []
|
|
|
|
def sample(self, batch_size):
|
|
if batch_size > 0:
|
|
indice = np.random.choice(self._size, batch_size)
|
|
else:
|
|
indice = np.arange(self._size)
|
|
return Batch(
|
|
obs=self.obs[indice],
|
|
act=self.act[indice],
|
|
rew=self.rew[indice],
|
|
done=self.done[indice],
|
|
obs_next=self.obs_next[indice],
|
|
info=self.info[indice]
|
|
), indice
|
|
|
|
def __getitem__(self, index):
|
|
return Batch(
|
|
obs=self.obs[index],
|
|
act=self.act[index],
|
|
rew=self.rew[index],
|
|
done=self.done[index],
|
|
obs_next=self.obs_next[index],
|
|
info=self.info[index]
|
|
)
|
|
|
|
|
|
class PrioritizedReplayBuffer(ReplayBuffer):
|
|
"""docstring for PrioritizedReplayBuffer"""
|
|
|
|
def __init__(self, size):
|
|
super().__init__(size)
|
|
|
|
def add(self, obs, act, rew, done, obs_next, info={}, weight=None):
|
|
raise NotImplementedError
|
|
|
|
def sample_indice(self, batch_size):
|
|
raise NotImplementedError
|
|
|
|
def sample(self, batch_size):
|
|
raise NotImplementedError
|