-
Notifications
You must be signed in to change notification settings - Fork 2
/
tree.py
135 lines (107 loc) · 4.3 KB
/
tree.py
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import numpy as np
class SumTree(object):
"""
This SumTree code is modified version of Morvan Zhou:
https://github.com/MorvanZhou/Reinforcement-learning-with-tensorflow/blob/master/contents/5.2_Prioritized_Replay_DQN/RL_brain.py
"""
data_pointer = 0
"""
Here we initialize the tree with all nodes = 0, and initialize the data with all values = 0
"""
def __init__(self, capacity):
self.capacity = capacity # Number of leaf nodes (final nodes) that contains experiences
# Generate the tree with all nodes values = 0
# To understand this calculation (2 * capacity - 1) look at the schema above
# Remember we are in a binary node (each node has max 2 children) so 2x size of leaf (capacity) - 1 (root node)
# Parent nodes = capacity - 1
# Leaf nodes = capacity
self.tree = np.zeros(2 * capacity - 1)
""" tree:
0
/ \
0 0
/ \ / \
0 0 0 0 [Size: capacity] it's at this line that there is the priorities score (aka pi)
"""
# Contains the experiences (so the size of data is capacity)
self.data = np.zeros(capacity, dtype=object)
"""
Here we add our priority score in the sumtree leaf and add the experience in data
"""
def add(self, priority, data):
# Look at what index we want to put the experience
tree_index = self.data_pointer + self.capacity - 1
""" tree:
0
/ \
0 0
/ \ / \
tree_index 0 0 0 We fill the leaves from left to right
"""
# Update data frame
self.data[self.data_pointer] = data
# Update the leaf
self.update(tree_index, priority)
# Add 1 to data_pointer
self.data_pointer += 1
if self.data_pointer >= self.capacity: # If we're above the capacity, you go back to first index (we overwrite)
self.data_pointer = 0
"""
Update the leaf priority score and propagate the change through tree
"""
def update(self, tree_index, priority):
# Change = new priority score - former priority score
change = priority - self.tree[tree_index]
self.tree[tree_index] = priority
# then propagate the change through tree
while tree_index != 0: # this method is faster than the recursive loop in the reference code
"""
Here we want to access the line above
THE NUMBERS IN THIS TREE ARE THE INDEXES NOT THE PRIORITY VALUES
0
/ \
1 2
/ \ / \
3 4 5 [6]
If we are in leaf at index 6, we updated the priority score
We need then to update index 2 node
So tree_index = (tree_index - 1) // 2
tree_index = (6-1)//2
tree_index = 2 (because // round the result)
"""
tree_index = (tree_index - 1) // 2
self.tree[tree_index] += change
"""
Here we get the leaf_index, priority value of that leaf and experience associated with that index
"""
def get_leaf(self, v):
"""
Tree structure and array storage:
Tree index:
0 -> storing priority sum
/ \
1 2
/ \ / \
3 4 5 6 -> storing priority for experiences
Array type for storing:
[0,1,2,3,4,5,6]
"""
parent_index = 0
while True: # the while loop is faster than the method in the reference code
left_child_index = 2 * parent_index + 1
right_child_index = left_child_index + 1
# If we reach bottom, end the search
if left_child_index >= len(self.tree):
leaf_index = parent_index
break
else: # downward search, always search for a higher priority node
if v <= self.tree[left_child_index]:
parent_index = left_child_index
else:
v -= self.tree[left_child_index]
parent_index = right_child_index
data_index = leaf_index - self.capacity + 1
return leaf_index, self.tree[leaf_index], self.data[data_index]
@property
def total_priority(self):
return self.tree[0] # Returns the root node