prioritized_experience_replay.py

#
# Copyright (c) 2017 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import operator
import random
from enum import Enum
from typing import List, Tuple, Any

import numpy as np

from rl_coach.core_types import Transition
from rl_coach.memories.memory import MemoryGranularity
from rl_coach.memories.non_episodic.experience_replay import ExperienceReplayParameters, ExperienceReplay
from rl_coach.schedules import Schedule, ConstantSchedule


class PrioritizedExperienceReplayParameters(ExperienceReplayParameters):
    def __init__(self):
        super().__init__()
        self.max_size = (MemoryGranularity.Transitions, 1000000)
        self.alpha = 0.6
        self.beta = ConstantSchedule(0.4)
        self.epsilon = 1e-6

    @property
    def path(self):
        return 'rl_coach.memories.non_episodic.prioritized_experience_replay:PrioritizedExperienceReplay'


class SegmentTree(object):
    """
    A tree which can be used as a min/max heap or a sum tree
    Add or update item value - O(log N)
    Sampling an item - O(log N)
    """
    class Operation(Enum):
        MAX = {"operator": max, "initial_value": -float("inf")}
        MIN = {"operator": min, "initial_value": float("inf")}
        SUM = {"operator": operator.add, "initial_value": 0}

    def __init__(self, size: int, operation: Operation):
        self.next_leaf_idx_to_write = 0
        self.size = size
        if not (size > 0 and size & (size - 1) == 0):
            raise ValueError("A segment tree size must be a positive power of 2. The given size is {}".format(self.size))
        self.operation = operation
        self.tree = np.ones(2 * size - 1) * self.operation.value['initial_value']
        self.data = [None] * size

    def _propagate(self, node_idx: int) -> None:
        """
        Propagate an update of a node's value to its parent node
        :param node_idx: the index of the node that was updated
        :return: None
        """
        parent = (node_idx - 1) // 2

        self.tree[parent] = self.operation.value['operator'](self.tree[parent * 2 + 1], self.tree[parent * 2 + 2])

        if parent != 0:
            self._propagate(parent)

    def _retrieve(self, root_node_idx: int, val: float)-> int:
        """
        Retrieve the first node that has a value larger than val and is a child of the node at index idx
        :param root_node_idx: the index of the root node to search from
        :param val: the value to query for
        :return: the index of the resulting node
        """
        left = 2 * root_node_idx + 1
        right = left + 1

        if left >= len(self.tree):
            return root_node_idx

        if val <= self.tree[left]:
            return self._retrieve(left, val)
        else:
            return self._retrieve(right, val-self.tree[left])

    def total_value(self) -> float:
        """
        Return the total value of the tree according to the tree operation. For SUM for example, this will return
        the total sum of the tree. for MIN, this will return the minimal value
        :return: the total value of the tree
        """
        return self.tree[0]

    def add(self, val: float, data: Any) -> None:
        """
        Add a new value to the tree with data assigned to it
        :param val: the new value to add to the tree
        :param data: the data that should be assigned to this value
        :return: None
        """
        self.data[self.next_leaf_idx_to_write] = data
        self.update(self.next_leaf_idx_to_write, val)

        self.next_leaf_idx_to_write += 1
        if self.next_leaf_idx_to_write >= self.size:
            self.next_leaf_idx_to_write = 0

    def update(self, leaf_idx: int, new_val: float) -> None:
        """
        Update the value of the node at index idx
        :param leaf_idx: the index of the node to update
        :param new_val: the new value of the node
        :return: None
        """
        node_idx = leaf_idx + self.size - 1
        if not 0 <= node_idx < len(self.tree):
            raise ValueError("The given left index ({}) can not be found in the tree. The available leaves are: 0-{}"
                             .format(leaf_idx, self.size - 1))

        self.tree[node_idx] = new_val
        self._propagate(node_idx)

    def get_element_by_partial_sum(self, val: float) -> Tuple[int, float, Any]:
        """
        Given a value between 0 and the tree sum, return the object which this value is in it's range.
        For example, if we have 3 leaves: 10, 20, 30, and val=35, this will return the 3rd leaf, by accumulating
        leaves by their order until getting to 35. This allows sampling leaves according to their proportional
        probability.
        :param val: a value within the range 0 and the tree sum
        :return: the index of the resulting leaf in the tree, its probability and
                 the object itself
        """
        node_idx = self._retrieve(0, val)
        leaf_idx = node_idx - self.size + 1
        data_value = self.tree[node_idx]
        data = self.data[leaf_idx]

        return leaf_idx, data_value, data

    def __str__(self):
        result = ""
        start = 0
        size = 1
        while size <= self.size:
            result += "{}\n".format(self.tree[start:(start + size)])
            start += size
            size *= 2
        return result


class PrioritizedExperienceReplay(ExperienceReplay):
    """
    This is the proportional sampling variant of the prioritized experience replay as described
    in https://arxiv.org/pdf/1511.05952.pdf.
    """
    def __init__(self, max_size: Tuple[MemoryGranularity, int], alpha: float=0.6, beta: Schedule=ConstantSchedule(0.4),
                 epsilon: float=1e-6, allow_duplicates_in_batch_sampling: bool=True):
        """
        :param max_size: the maximum number of transitions or episodes to hold in the memory
        :param alpha: the alpha prioritization coefficient
        :param beta: the beta parameter used for importance sampling
        :param epsilon: a small value added to the priority of each transition
        :param allow_duplicates_in_batch_sampling: allow having the same transition multiple times in a batch
        """
        if max_size[0] != MemoryGranularity.Transitions:
            raise ValueError("Prioritized Experience Replay currently only support setting the memory size in "
                             "transitions granularity.")
        self.power_of_2_size = 1
        while self.power_of_2_size < max_size[1]:
            self.power_of_2_size *= 2
        super().__init__((MemoryGranularity.Transitions, self.power_of_2_size), allow_duplicates_in_batch_sampling)
        self.sum_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.SUM)
        self.min_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MIN)
        self.max_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MAX)
        self.alpha = alpha
        self.beta = beta
        self.epsilon = epsilon
        self.maximal_priority = 1.0

    def _update_priority(self, leaf_idx: int, error: float) -> None:
        """
        Update the priority of a given transition, using its index in the tree and its error
        :param leaf_idx: the index of the transition leaf in the tree
        :param error: the new error value
        :return: None
        """
        if error < 0:
            raise ValueError("The priorities must be non-negative values")
        priority = (error + self.epsilon)
        self.sum_tree.update(leaf_idx, priority ** self.alpha)
        self.min_tree.update(leaf_idx, priority ** self.alpha)
        self.max_tree.update(leaf_idx, priority)
        self.maximal_priority = self.max_tree.total_value()

    def update_priorities(self, indices: List[int], error_values: List[float]) -> None:
        """
        Update the priorities of a batch of transitions using their indices and their new TD error terms
        :param indices: the indices of the transitions to update
        :param error_values: the new error values
        :return: None
        """
        self.reader_writer_lock.lock_writing_and_reading()

        if len(indices) != len(error_values):
            raise ValueError("The number of indexes requested for update don't match the number of error values given")
        for transition_idx, error in zip(indices, error_values):
            self._update_priority(transition_idx, error)

        self.reader_writer_lock.release_writing_and_reading()

    def sample(self, size: int) -> List[Transition]:
        """
        Sample a batch of transitions form the replay buffer. If the requested size is larger than the number
        of samples available in the replay buffer then the batch will return empty.
        :param size: the size of the batch to sample
        :return: a batch (list) of selected transitions from the replay buffer
        """

        self.reader_writer_lock.lock_writing()

        if self.num_transitions() >= size:
            # split the tree leaves to equal segments and sample one transition from each segment
            batch = []
            segment_size = self.sum_tree.total_value() / size

            # get the maximum weight in the memory
            min_probability = self.min_tree.total_value() / self.sum_tree.total_value()  # min P(j) = min p^a / sum(p^a)
            max_weight = (min_probability * self.num_transitions()) ** -self.beta.current_value  # max wi

            # sample a batch
            for i in range(size):
                segment_start = segment_size * i
                segment_end = segment_size * (i + 1)

                # sample leaf and calculate its weight
                val = random.uniform(segment_start, segment_end)
                leaf_idx, priority, transition = self.sum_tree.get_element_by_partial_sum(val)
                priority /= self.sum_tree.total_value()   # P(j) = p^a / sum(p^a)
                weight = (self.num_transitions() * priority) ** -self.beta.current_value  # (N * P(j)) ^ -beta
                normalized_weight = weight / max_weight  # wj = ((N * P(j)) ^ -beta) / max wi

                transition.info['idx'] = leaf_idx
                transition.info['weight'] = normalized_weight

                batch.append(transition)

            self.beta.step()

        else:
            raise ValueError("The replay buffer cannot be sampled since there are not enough transitions yet. "
                             "There are currently {} transitions".format(self.num_transitions()))

        self.reader_writer_lock.release_writing()
        return batch

    def store(self, transition: Transition, lock=True) -> None:
        """
        Store a new transition in the memory.
        :param transition: a transition to store
        :return: None
        """
        # Calling super.store() so that in case a memory backend is used, the memory backend can store this transition.
        super().store(transition)

        if lock:
            self.reader_writer_lock.lock_writing_and_reading()

        transition_priority = self.maximal_priority
        self.sum_tree.add(transition_priority ** self.alpha, transition)
        self.min_tree.add(transition_priority ** self.alpha, transition)
        self.max_tree.add(transition_priority, transition)
        super().store(transition, False)

        if lock:
            self.reader_writer_lock.release_writing_and_reading()

    def clean(self, lock=True) -> None:
        """
        Clean the memory by removing all the episodes
        :return: None
        """
        if lock:
            self.reader_writer_lock.lock_writing_and_reading()

        super().clean(lock=False)
        self.sum_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.SUM)
        self.min_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MIN)
        self.max_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MAX)

        if lock:
            self.reader_writer_lock.release_writing_and_reading()