<reponame>schroeder-dewitt/maddpg<gh_stars>1-10
import numpy as np
import random
# from baselines.common.segment_tree import SumSegmentTree, MinSegmentTree
from maddpg.common.segment_tree import SumSegmentTree, MinSegmentTree

class _ReplayBuffer(object):
    def __init__(self, size, experience_size, q_lstm_on, p_lstm_on):
        """Create Prioritized Replay buffer.
        Parameters
        ----------
        size: int
            Max number of transitions to store in the buffer. When the buffer
            overflows the old memories are dropped.
        """
        self._storage = []
        self._maxsize = int(size)
        self._next_idx = 0
        self.experience_size = experience_size
        self.first_ever_episode = True
        self.q_lstm_on = q_lstm_on
        self.p_lstm_on = p_lstm_on

    def __len__(self):
        return len(self._storage)

    def clear(self):
        self._storage = []
        self._next_idx = 0

    def add(self, *args ):
        assert len(args) == self.experience_size, "[Replay Buffer] Got exp size {}, but expected size {}".format(len(args),self.experience_size)
        # buffer: [[(episode exp, ...), (episode exp),... ], [(diff ep...), (...)]]
        # data = (obs_t, action, reward, obs_tp1, done,
        #             p_c_in, p_h_in, p_c_out, p_h_out,
        #             q_c_in, q_h_in, q_c_out, q_h_out)
        data = args[:-1]
        new_episode = args[-1]
        if new_episode:
            if self.first_ever_episode:
                self.first_ever_episode = False
            else:
                self._next_idx = (self._next_idx + 1) % self._maxsize

        if self._next_idx >= len(self._storage):
            # still have room for more episodes
            if new_episode:
                # start a new episode list @ storage[next_ind]
                self._storage.append([data])
            else:
                # add onto current episode length
                self._storage[self._next_idx].append(data)
        else: # replace old data
            if new_episode:
                self._storage[self._next_idx] = [data]
            else:
                self._storage[self._next_idx].append(data)

    def _encode_sample(self, idxes):
        obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
        p_c_in, p_h_in= [], []
        p_c_out, p_h_out = [], []
        q_c_in, q_h_in= [], []
        q_c_out, q_h_out = [], []
        for i in idxes:
            data = self._storage[i]
            for el in data:
                if self.q_lstm_on and self.p_lstm_on:
                    obs_t, action, reward, obs_tp1, done, p_c_in_t, p_h_in_t, p_c_out_t, p_h_out_t,q_c_in_t, q_h_in_t, q_c_out_t, q_h_out_t, state, state_tp1 = el
                elif self.q_lstm_on:
                    obs_t, action, reward, obs_tp1, done, q_c_in_t, q_h_in_t, q_c_out_t, q_h_out_t, state, state_tp1 = el
                elif self.p_lstm_on:
                    obs_t, action, reward, obs_tp1, done, p_c_in_t, p_h_in_t, p_c_out_t, p_h_out_t, state, state_tp1 = el
                else:
                    obs_t, action, reward, obs_tp1, done, state, state_tp1  = el

                # might need to change the dimensions on obs and actions
                obses_t.append(np.array(obs_t, copy=False))
                actions.append(np.array(action, copy=False))
                rewards.append(reward)
                obses_tp1.append(np.array(obs_tp1, copy=False))
                dones.append(done)

                if self.p_lstm_on:
                    p_c_in.append(p_c_in_t)
                    p_h_in.append(p_h_in_t)
                    p_c_out.append(p_c_out_t)
                    p_h_out.append(p_h_out_t)
                if self.q_lstm_on:
                    q_c_in.append(q_c_in_t)
                    q_h_in.append(q_h_in_t)
                    q_c_out.append(q_c_out_t)
                    q_h_out.append(q_h_out_t)
            # Check the dimensions
        return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones), np.array(p_c_in),np.array(p_h_in),np.array(p_c_out),np.array(p_h_out), np.array(q_c_in),np.array(q_h_in),np.array(q_c_out),np.array(q_h_out), np.array(state), np.array(state_tp1)


    def make_index_lstm(self, batch_size):
        res = []
        for _ in range(batch_size):
            rand_int = random.randint(0, len(self._storage) -1)
            res.append(rand_int)
        return res

    def make_latest_index(self, batch_size):
        idx = [(self._next_idx - 1 - i) % self._maxsize for i in range(batch_size)]
        np.random.shuffle(idx)
        return idx

    def sample_index(self, idxes):
        return self._encode_sample(idxes)

    def sample(self, batch_size):
        """Sample a batch of experiences.
        Parameters
        ----------
        batch_size: int
            How many transitions to sample.
        Returns
        -------
        obs_batch: np.array
            batch of observations
        act_batch: np.array
            batch of actions executed given obs_batch
        rew_batch: np.array
            rewards received as results of executing act_batch
        next_obs_batch: np.array
            next set of observations seen after executing act_batch
        done_mask: np.array
            done_mask[i] = 1 if executing act_batch[i] resulted in
            the end of an episode and 0 otherwise.
        """
        if batch_size > 0:
            idxes = self.make_index(batch_size)
        else:
            idxes = range(0, len(self._storage))
        return self._encode_sample(idxes)

    def collect(self):
        return self.sample(-1)


class ReplayBufferLSTM(object):
    def __init__(self, size):
        """Create Prioritized Replay buffer.
        Parameters
        ----------
        size: int
            Max number of transitions to store in the buffer. When the buffer
            overflows the old memories are dropped.
        """
        self._storage = []
        self._maxsize = int(size)
        self._next_idx = 0

        self.first_ever_episode = True

    def __len__(self):
        return len(self._storage)

    def clear(self):
        self._storage = []
        self._next_idx = 0

    def add(self, obs_t, action, reward, obs_tp1, done,
                    p_c_in, p_h_in, p_c_out, p_h_out,
                    q_c_in, q_h_in, q_c_out, q_h_out,  new_episode=False):
        # buffer: [[(episode exp, ...), (episode exp),... ], [(diff ep...), (...)]]
        data = (obs_t, action, reward, obs_tp1, done,
                    p_c_in, p_h_in, p_c_out, p_h_out,
                    q_c_in, q_h_in, q_c_out, q_h_out)
        if new_episode:
            if self.first_ever_episode:
                self.first_ever_episode = False
            else:
                self._next_idx = (self._next_idx + 1) % self._maxsize

        if self._next_idx >= len(self._storage):
            # still have room for more episodes
            if new_episode:
                # start a new episode list @ storage[next_ind]
                self._storage.append([data])
            else:
                # add onto current episode length
                self._storage[self._next_idx].append(data)
        else: # replace old data
            if new_episode:
                self._storage[self._next_idx] = [data]
            else:
                self._storage[self._next_idx].append(data)

    def _encode_sample(self, idxes):
        obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
        p_c_in, p_h_in= [], []
        p_c_out, p_h_out = [], []
        q_c_in, q_h_in= [], []
        q_c_out, q_h_out = [], []
        for i in idxes:
            data = self._storage[i]
            for el in data:
                obs_t, action, reward, obs_tp1, done, p_c_in_t, p_h_in_t, p_c_out_t, p_h_out_t,q_c_in_t, q_h_in_t, q_c_out_t, q_h_out_t= el
                # might need to change the dimensions on obs and actions
                obses_t.append(np.array(obs_t, copy=False))
                actions.append(np.array(action, copy=False))
                rewards.append(reward)
                obses_tp1.append(np.array(obs_tp1, copy=False))
                dones.append(done)

                p_c_in.append(p_c_in_t)
                p_h_in.append(p_h_in_t)
                p_c_out.append(p_c_out_t)
                p_h_out.append(p_h_out_t)

                q_c_in.append(q_c_in_t)
                q_h_in.append(q_h_in_t)
                q_c_out.append(q_c_out_t)
                q_h_out.append(q_h_out_t)
            # Check the dimensions
        return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones), np.array(p_c_in),np.array(p_h_in),np.array(p_c_out),np.array(p_h_out), np.array(q_c_in),np.array(q_h_in),np.array(q_c_out),np.array(q_h_out)


    def make_index_lstm(self, batch_size):
        # return [random.randint(0, len(self._storage) - 1) for _ in range(batch_size)]
        res = []
        for _ in range(batch_size):
            rand_int = random.randint(0, len(self._storage) -1)
            # for i in range(20):
            #    res += [rand_int + i]
            #rand_int = random.randint(0, len(self._storage) -1)
            res.append(rand_int)
        return res

    def make_latest_index(self, batch_size):
        idx = [(self._next_idx - 1 - i) % self._maxsize for i in range(batch_size)]
        np.random.shuffle(idx)
        return idx

    def sample_index(self, idxes):
        return self._encode_sample(idxes)

    def sample(self, batch_size):
        """Sample a batch of experiences.
        Parameters
        ----------
        batch_size: int
            How many transitions to sample.
        Returns
        -------
        obs_batch: np.array
            batch of observations
        act_batch: np.array
            batch of actions executed given obs_batch
        rew_batch: np.array
            rewards received as results of executing act_batch
        next_obs_batch: np.array
            next set of observations seen after executing act_batch
        done_mask: np.array
            done_mask[i] = 1 if executing act_batch[i] resulted in
            the end of an episode and 0 otherwise.
        """
        if batch_size > 0:
            idxes = self.make_index(batch_size)
        else:
            idxes = range(0, len(self._storage))
        return self._encode_sample(idxes)

    def collect(self):
        return self.sample(-1)



class PrioritizedReplayBuffer(ReplayBufferLSTM):
    def __init__(self, size, alpha):
        """Create Prioritized Replay buffer.
        Parameters
        ----------
        size: int
            Max number of transitions to store in the buffer. When the buffer
            overflows the old memories are dropped.
        alpha: float
            how much prioritization is used
            (0 - no prioritization, 1 - full prioritization)
        See Also
        --------
        ReplayBuffer.__init__
        """
        super(PrioritizedReplayBuffer, self).__init__(size)
        assert alpha >= 0
        self._alpha = alpha

        it_capacity = 1
        while it_capacity < size:
            it_capacity *= 2

        self._it_sum = SumSegmentTree(it_capacity)
        self._it_min = MinSegmentTree(it_capacity)
        self._max_priority = 1.0

    def add(self, *args, **kwargs):
        """See ReplayBuffer.store_effect"""
        idx = self._next_idx
        super().add(*args, **kwargs)
        self._it_sum[idx] = self._max_priority ** self._alpha
        self._it_min[idx] = self._max_priority ** self._alpha

    def _sample_proportional(self, batch_size):
        res = []
        p_total = self._it_sum.sum(0, len(self._storage) - 1)
        every_range_len = p_total / batch_size
        for i in range(batch_size):
            mass = random.random() * every_range_len + i * every_range_len
            idx = self._it_sum.find_prefixsum_idx(mass)
            res.append(idx)
        return res

    def sample(self, batch_size, beta):
        """Sample a batch of experiences.
        compared to ReplayBuffer.sample
        it also returns importance weights and idxes
        of sampled experiences.
        Parameters
        ----------
        batch_size: int
            How many transitions to sample.
        beta: float
            To what degree to use importance weights
            (0 - no corrections, 1 - full correction)
        Returns
        -------
        obs_batch: np.array
            batch of observations
        act_batch: np.array
            batch of actions executed given obs_batch
        rew_batch: np.array
            rewards received as results of executing act_batch
        next_obs_batch: np.array
            next set of observations seen after executing act_batch
        done_mask: np.array
            done_mask[i] = 1 if executing act_batch[i] resulted in
            the end of an episode and 0 otherwise.
        weights: np.array
            Array of shape (batch_size,) and dtype np.float32
            denoting importance weight of each sampled transition
        idxes: np.array
            Array of shape (batch_size,) and dtype np.int32
            idexes in buffer of sampled experiences
        """
        assert beta > 0

        idxes = self._sample_proportional(batch_size)

        weights = []
        p_min = self._it_min.min() / self._it_sum.sum()
        max_weight = (p_min * len(self._storage)) ** (-beta)

        for idx in idxes:
            p_sample = self._it_sum[idx] / self._it_sum.sum()
            weight = (p_sample * len(self._storage)) ** (-beta)
            weights.append(weight / max_weight)
        weights = np.array(weights)
        encoded_sample = self._encode_sample(idxes)
        return tuple(list(encoded_sample) + [weights, idxes])

    def update_priorities(self, idxes, priorities):
        """Update priorities of sampled transitions.
        sets priority of transition at index idxes[i] in buffer
        to priorities[i].
        Parameters
        ----------
        idxes: [int]
            List of idxes of sampled transitions
        priorities: [float]
            List of updated priorities corresponding to
            transitions at the sampled idxes denoted by
            variable `idxes`.
        """
        assert len(idxes) == len(priorities)
        for idx, priority in zip(idxes, priorities):
            assert priority > 0
            assert 0 <= idx < len(self._storage)
            self._it_sum[idx] = priority ** self._alpha
            self._it_min[idx] = priority ** self._alpha

            self._max_priority = max(self._max_priority, priority)

import numpy as np
import random
# from baselines.common.segment_tree import SumSegmentTree, MinSegmentTree
from maddpg.common.segment_tree import SumSegmentTree, MinSegmentTree

class _ReplayBuffer(object):
    def __init__(self, size, experience_size, q_lstm_on, p_lstm_on):
        """Create Prioritized Replay buffer.
        Parameters
        ----------
        size: int
            Max number of transitions to store in the buffer. When the buffer
            overflows the old memories are dropped.
        """
        self._storage = []
        self._maxsize = int(size)
        self._next_idx = 0
        self.experience_size = experience_size
        self.first_ever_episode = True
        self.q_lstm_on = q_lstm_on
        self.p_lstm_on = p_lstm_on

    def __len__(self):
        return len(self._storage)

    def clear(self):
        self._storage = []
        self._next_idx = 0

    def add(self, *args ):
        assert len(args) == self.experience_size, "[Replay Buffer] Got exp size {}, but expected size {}".format(len(args),self.experience_size)
        # buffer: [[(episode exp, ...), (episode exp),... ], [(diff ep...), (...)]]
        # data = (obs_t, action, reward, obs_tp1, done,
        #             p_c_in, p_h_in, p_c_out, p_h_out,
        #             q_c_in, q_h_in, q_c_out, q_h_out)
        data = args[:-1]
        new_episode = args[-1]
        if new_episode:
            if self.first_ever_episode:
                self.first_ever_episode = False
            else:
                self._next_idx = (self._next_idx + 1) % self._maxsize

        if self._next_idx >= len(self._storage):
            # still have room for more episodes
            if new_episode:
                # start a new episode list @ storage[next_ind]
                self._storage.append([data])
            else:
                # add onto current episode length
                self._storage[self._next_idx].append(data)
        else: # replace old data
            if new_episode:
                self._storage[self._next_idx] = [data]
            else:
                self._storage[self._next_idx].append(data)

    def _encode_sample(self, idxes):
        obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
        p_c_in, p_h_in= [], []
        p_c_out, p_h_out = [], []
        q_c_in, q_h_in= [], []
        q_c_out, q_h_out = [], []
        for i in idxes:
            data = self._storage[i]
            for el in data:
                if self.q_lstm_on and self.p_lstm_on:
                    obs_t, action, reward, obs_tp1, done, p_c_in_t, p_h_in_t, p_c_out_t, p_h_out_t,q_c_in_t, q_h_in_t, q_c_out_t, q_h_out_t, state, state_tp1 = el
                elif self.q_lstm_on:
                    obs_t, action, reward, obs_tp1, done, q_c_in_t, q_h_in_t, q_c_out_t, q_h_out_t, state, state_tp1 = el
                elif self.p_lstm_on:
                    obs_t, action, reward, obs_tp1, done, p_c_in_t, p_h_in_t, p_c_out_t, p_h_out_t, state, state_tp1 = el
                else:
                    obs_t, action, reward, obs_tp1, done, state, state_tp1  = el

                # might need to change the dimensions on obs and actions
                obses_t.append(np.array(obs_t, copy=False))
                actions.append(np.array(action, copy=False))
                rewards.append(reward)
                obses_tp1.append(np.array(obs_tp1, copy=False))
                dones.append(done)

                if self.p_lstm_on:
                    p_c_in.append(p_c_in_t)
                    p_h_in.append(p_h_in_t)
                    p_c_out.append(p_c_out_t)
                    p_h_out.append(p_h_out_t)
                if self.q_lstm_on:
                    q_c_in.append(q_c_in_t)
                    q_h_in.append(q_h_in_t)
                    q_c_out.append(q_c_out_t)
                    q_h_out.append(q_h_out_t)
            # Check the dimensions
        return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones), np.array(p_c_in),np.array(p_h_in),np.array(p_c_out),np.array(p_h_out), np.array(q_c_in),np.array(q_h_in),np.array(q_c_out),np.array(q_h_out), np.array(state), np.array(state_tp1)


    def make_index_lstm(self, batch_size):
        res = []
        for _ in range(batch_size):
            rand_int = random.randint(0, len(self._storage) -1)
            res.append(rand_int)
        return res

    def make_latest_index(self, batch_size):
        idx = [(self._next_idx - 1 - i) % self._maxsize for i in range(batch_size)]
        np.random.shuffle(idx)
        return idx

    def sample_index(self, idxes):
        return self._encode_sample(idxes)

    def sample(self, batch_size):
        """Sample a batch of experiences.
        Parameters
        ----------
        batch_size: int
            How many transitions to sample.
        Returns
        -------
        obs_batch: np.array
            batch of observations
        act_batch: np.array
            batch of actions executed given obs_batch
        rew_batch: np.array
            rewards received as results of executing act_batch
        next_obs_batch: np.array
            next set of observations seen after executing act_batch
        done_mask: np.array
            done_mask[i] = 1 if executing act_batch[i] resulted in
            the end of an episode and 0 otherwise.
        """
        if batch_size > 0:
            idxes = self.make_index(batch_size)
        else:
            idxes = range(0, len(self._storage))
        return self._encode_sample(idxes)

    def collect(self):
        return self.sample(-1)


class ReplayBufferLSTM(object):
    def __init__(self, size):
        """Create Prioritized Replay buffer.
        Parameters
        ----------
        size: int
            Max number of transitions to store in the buffer. When the buffer
            overflows the old memories are dropped.
        """
        self._storage = []
        self._maxsize = int(size)
        self._next_idx = 0

        self.first_ever_episode = True

    def __len__(self):
        return len(self._storage)

    def clear(self):
        self._storage = []
        self._next_idx = 0

    def add(self, obs_t, action, reward, obs_tp1, done,
                    p_c_in, p_h_in, p_c_out, p_h_out,
                    q_c_in, q_h_in, q_c_out, q_h_out,  new_episode=False):
        # buffer: [[(episode exp, ...), (episode exp),... ], [(diff ep...), (...)]]
        data = (obs_t, action, reward, obs_tp1, done,
                    p_c_in, p_h_in, p_c_out, p_h_out,
                    q_c_in, q_h_in, q_c_out, q_h_out)
        if new_episode:
            if self.first_ever_episode:
                self.first_ever_episode = False
            else:
                self._next_idx = (self._next_idx + 1) % self._maxsize

        if self._next_idx >= len(self._storage):
            # still have room for more episodes
            if new_episode:
                # start a new episode list @ storage[next_ind]
                self._storage.append([data])
            else:
                # add onto current episode length
                self._storage[self._next_idx].append(data)
        else: # replace old data
            if new_episode:
                self._storage[self._next_idx] = [data]
            else:
                self._storage[self._next_idx].append(data)

    def _encode_sample(self, idxes):
        obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
        p_c_in, p_h_in= [], []
        p_c_out, p_h_out = [], []
        q_c_in, q_h_in= [], []
        q_c_out, q_h_out = [], []
        for i in idxes:
            data = self._storage[i]
            for el in data:
                obs_t, action, reward, obs_tp1, done, p_c_in_t, p_h_in_t, p_c_out_t, p_h_out_t,q_c_in_t, q_h_in_t, q_c_out_t, q_h_out_t= el
                # might need to change the dimensions on obs and actions
                obses_t.append(np.array(obs_t, copy=False))
                actions.append(np.array(action, copy=False))
                rewards.append(reward)
                obses_tp1.append(np.array(obs_tp1, copy=False))
                dones.append(done)

                p_c_in.append(p_c_in_t)
                p_h_in.append(p_h_in_t)
                p_c_out.append(p_c_out_t)
                p_h_out.append(p_h_out_t)

                q_c_in.append(q_c_in_t)
                q_h_in.append(q_h_in_t)
                q_c_out.append(q_c_out_t)
                q_h_out.append(q_h_out_t)
            # Check the dimensions
        return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones), np.array(p_c_in),np.array(p_h_in),np.array(p_c_out),np.array(p_h_out), np.array(q_c_in),np.array(q_h_in),np.array(q_c_out),np.array(q_h_out)


    def make_index_lstm(self, batch_size):
        # return [random.randint(0, len(self._storage) - 1) for _ in range(batch_size)]
        res = []
        for _ in range(batch_size):
            rand_int = random.randint(0, len(self._storage) -1)
            # for i in range(20):
            #    res += [rand_int + i]
            #rand_int = random.randint(0, len(self._storage) -1)
            res.append(rand_int)
        return res

    def make_latest_index(self, batch_size):
        idx = [(self._next_idx - 1 - i) % self._maxsize for i in range(batch_size)]
        np.random.shuffle(idx)
        return idx

    def sample_index(self, idxes):
        return self._encode_sample(idxes)

    def sample(self, batch_size):
        """Sample a batch of experiences.
        Parameters
        ----------
        batch_size: int
            How many transitions to sample.
        Returns
        -------
        obs_batch: np.array
            batch of observations
        act_batch: np.array
            batch of actions executed given obs_batch
        rew_batch: np.array
            rewards received as results of executing act_batch
        next_obs_batch: np.array
            next set of observations seen after executing act_batch
        done_mask: np.array
            done_mask[i] = 1 if executing act_batch[i] resulted in
            the end of an episode and 0 otherwise.
        """
        if batch_size > 0:
            idxes = self.make_index(batch_size)
        else:
            idxes = range(0, len(self._storage))
        return self._encode_sample(idxes)

    def collect(self):
        return self.sample(-1)



class PrioritizedReplayBuffer(ReplayBufferLSTM):
    def __init__(self, size, alpha):
        """Create Prioritized Replay buffer.
        Parameters
        ----------
        size: int
            Max number of transitions to store in the buffer. When the buffer
            overflows the old memories are dropped.
        alpha: float
            how much prioritization is used
            (0 - no prioritization, 1 - full prioritization)
        See Also
        --------
        ReplayBuffer.__init__
        """
        super(PrioritizedReplayBuffer, self).__init__(size)
        assert alpha >= 0
        self._alpha = alpha

        it_capacity = 1
        while it_capacity < size:
            it_capacity *= 2

        self._it_sum = SumSegmentTree(it_capacity)
        self._it_min = MinSegmentTree(it_capacity)
        self._max_priority = 1.0

    def add(self, *args, **kwargs):
        """See ReplayBuffer.store_effect"""
        idx = self._next_idx
        super().add(*args, **kwargs)
        self._it_sum[idx] = self._max_priority ** self._alpha
        self._it_min[idx] = self._max_priority ** self._alpha

    def _sample_proportional(self, batch_size):
        res = []
        p_total = self._it_sum.sum(0, len(self._storage) - 1)
        every_range_len = p_total / batch_size
        for i in range(batch_size):
            mass = random.random() * every_range_len + i * every_range_len
            idx = self._it_sum.find_prefixsum_idx(mass)
            res.append(idx)
        return res

    def sample(self, batch_size, beta):
        """Sample a batch of experiences.
        compared to ReplayBuffer.sample
        it also returns importance weights and idxes
        of sampled experiences.
        Parameters
        ----------
        batch_size: int
            How many transitions to sample.
        beta: float
            To what degree to use importance weights
            (0 - no corrections, 1 - full correction)
        Returns
        -------
        obs_batch: np.array
            batch of observations
        act_batch: np.array
            batch of actions executed given obs_batch
        rew_batch: np.array
            rewards received as results of executing act_batch
        next_obs_batch: np.array
            next set of observations seen after executing act_batch
        done_mask: np.array
            done_mask[i] = 1 if executing act_batch[i] resulted in
            the end of an episode and 0 otherwise.
        weights: np.array
            Array of shape (batch_size,) and dtype np.float32
            denoting importance weight of each sampled transition
        idxes: np.array
            Array of shape (batch_size,) and dtype np.int32
            idexes in buffer of sampled experiences
        """
        assert beta > 0

        idxes = self._sample_proportional(batch_size)

        weights = []
        p_min = self._it_min.min() / self._it_sum.sum()
        max_weight = (p_min * len(self._storage)) ** (-beta)

        for idx in idxes:
            p_sample = self._it_sum[idx] / self._it_sum.sum()
            weight = (p_sample * len(self._storage)) ** (-beta)
            weights.append(weight / max_weight)
        weights = np.array(weights)
        encoded_sample = self._encode_sample(idxes)
        return tuple(list(encoded_sample) + [weights, idxes])

    def update_priorities(self, idxes, priorities):
        """Update priorities of sampled transitions.
        sets priority of transition at index idxes[i] in buffer
        to priorities[i].
        Parameters
        ----------
        idxes: [int]
            List of idxes of sampled transitions
        priorities: [float]
            List of updated priorities corresponding to
            transitions at the sampled idxes denoted by
            variable `idxes`.
        """
        assert len(idxes) == len(priorities)
        for idx, priority in zip(idxes, priorities):
            assert priority > 0
            assert 0 <= idx < len(self._storage)
            self._it_sum[idx] = priority ** self._alpha
            self._it_min[idx] = priority ** self._alpha

            self._max_priority = max(self._max_priority, priority)
# from baselines.common.segment_tree import SumSegmentTree, MinSegmentTree Create Prioritized Replay buffer.         Parameters         ----------         size: int             Max number of transitions to store in the buffer. When the buffer             overflows the old memories are dropped. # buffer: [[(episode exp, ...), (episode exp),... ], [(diff ep...), (...)]] # data = (obs_t, action, reward, obs_tp1, done, #             p_c_in, p_h_in, p_c_out, p_h_out, #             q_c_in, q_h_in, q_c_out, q_h_out) # still have room for more episodes # start a new episode list @ storage[next_ind] # add onto current episode length # replace old data # might need to change the dimensions on obs and actions # Check the dimensions Sample a batch of experiences.         Parameters         ----------         batch_size: int             How many transitions to sample.         Returns         -------         obs_batch: np.array             batch of observations         act_batch: np.array             batch of actions executed given obs_batch         rew_batch: np.array             rewards received as results of executing act_batch         next_obs_batch: np.array             next set of observations seen after executing act_batch         done_mask: np.array             done_mask[i] = 1 if executing act_batch[i] resulted in             the end of an episode and 0 otherwise. Create Prioritized Replay buffer.         Parameters         ----------         size: int             Max number of transitions to store in the buffer. When the buffer             overflows the old memories are dropped. # buffer: [[(episode exp, ...), (episode exp),... ], [(diff ep...), (...)]] # still have room for more episodes # start a new episode list @ storage[next_ind] # add onto current episode length # replace old data # might need to change the dimensions on obs and actions # Check the dimensions # return [random.randint(0, len(self._storage) - 1) for _ in range(batch_size)] # for i in range(20): #    res += [rand_int + i] #rand_int = random.randint(0, len(self._storage) -1) Sample a batch of experiences.         Parameters         ----------         batch_size: int             How many transitions to sample.         Returns         -------         obs_batch: np.array             batch of observations         act_batch: np.array             batch of actions executed given obs_batch         rew_batch: np.array             rewards received as results of executing act_batch         next_obs_batch: np.array             next set of observations seen after executing act_batch         done_mask: np.array             done_mask[i] = 1 if executing act_batch[i] resulted in             the end of an episode and 0 otherwise. Create Prioritized Replay buffer.         Parameters         ----------         size: int             Max number of transitions to store in the buffer. When the buffer             overflows the old memories are dropped.         alpha: float             how much prioritization is used             (0 - no prioritization, 1 - full prioritization)         See Also         --------         ReplayBuffer.__init__ See ReplayBuffer.store_effect Sample a batch of experiences.         compared to ReplayBuffer.sample         it also returns importance weights and idxes         of sampled experiences.         Parameters         ----------         batch_size: int             How many transitions to sample.         beta: float             To what degree to use importance weights             (0 - no corrections, 1 - full correction)         Returns         -------         obs_batch: np.array             batch of observations         act_batch: np.array             batch of actions executed given obs_batch         rew_batch: np.array             rewards received as results of executing act_batch         next_obs_batch: np.array             next set of observations seen after executing act_batch         done_mask: np.array             done_mask[i] = 1 if executing act_batch[i] resulted in             the end of an episode and 0 otherwise.         weights: np.array             Array of shape (batch_size,) and dtype np.float32             denoting importance weight of each sampled transition         idxes: np.array             Array of shape (batch_size,) and dtype np.int32             idexes in buffer of sampled experiences Update priorities of sampled transitions.         sets priority of transition at index idxes[i] in buffer         to priorities[i].         Parameters         ----------         idxes: [int]             List of idxes of sampled transitions         priorities: [float]             List of updated priorities corresponding to             transitions at the sampled idxes denoted by             variable `idxes`.