class_docstr: ## Buffer for Prioritized Experience Replay
  
  [Prioritized experience replay](https://arxiv.org/abs/1511.05952)
   samples important transitions more frequently.
  The transitions are prioritized by the Temporal Difference error (td error), $\delta$.
  
  We sample transition $i$ with probability,
  $$P(i) = rac{p_i^lpha}{\sum_k p_k^lpha}$$
  where $lpha$ is a hyper-parameter that determines how much
  prioritization is used, with $lpha = 0$ corresponding to uniform case.
  $p_i$ is the priority.
  
  We use proportional prioritization $p_i = |\delta_i| + \epsilon$ where
  $\delta_i$ is the temporal difference for transition $i$.
  
  We correct the bias introduced by prioritized replay using
   importance-sampling (IS) weights
  $$w_i = igg(rac{1}{N} rac{1}{P(i)}igg)^eta$$ in the loss function.
  This fully compensates when $eta = 1$.
  We normalize weights by $rac{1}{\max_i w_i}$ for stability.
  Unbiased nature is most important towards the convergence at end of training.
  Therefore we increase $eta$ towards end of training.
  
  ### Binary Segment Tree
  We use a binary segment tree to efficiently calculate
  $\sum_k^i p_k^lpha$, the cumulative probability,
  which is needed to sample.
  We also use a binary segment tree to find $\min p_i^lpha$,
  which is needed for $rac{1}{\max_i w_i}$.
  We can also use a min-heap for this.
  Binary Segment Tree lets us calculate these in $\mathcal{O}(\log n)$
  time, which is way more efficient that the naive $\mathcal{O}(n)$
  approach.
  
  This is how a binary segment tree works for sum;
  it is similar for minimum.
  Let $x_i$ be the list of $N$ values we want to represent.
  Let $b_{i,j}$ be the $j^{\mathop{th}}$ node of the $i^{\mathop{th}}$ row
   in the binary tree.
  That is two children of node $b_{i,j}$ are $b_{i+1,2j}$ and $b_{i+1,2j + 1}$.
  
  The leaf nodes on row $D = \left\lceil {1 + \log_2 N} 
  ight
  ceil$
   will have values of $x$.
  Every node keeps the sum of the two child nodes.
  That is, the root node keeps the sum of the entire array of values.
  The left and right children of the root node keep
   the sum of the first half of the array and
   the sum of the second half of the array, respectively.
  And so on...
  
  $$b_{i,j} = \sum_{k = (j -1) * 2^{D - i} + 1}^{j * 2^{D - i}} x_k$$
  
  Number of nodes in row $i$,
  $$N_i = \left\lceil{rac{N}{D - i + 1}} 
  ight
  ceil$$
  This is equal to the sum of nodes in all rows above $i$.
  So we can use a single array $a$ to store the tree, where,
  $$b_{i,j} 
  ightarrow a_{N_i + j}$$
  
  Then child nodes of $a_i$ are $a_{2i}$ and $a_{2i + 1}$.
  That is,
  $$a_i = a_{2i} + a_{2i + 1}$$
  
  This way of maintaining binary trees is very easy to program.
  *Note that we are indexing starting from 1*.
  
  We use the same structure to compute the minimum.
class_docstr_tok: ["#","#","Buffer","for","Prioritized","Experience","Replay","[","Prioritized","experience","replay","]","(","https",":","\/\/arxiv.org\/abs\/1511.05952",")","samples","important","transitions","more","frequently",".","The","transitions","are","prioritized","by","the","Temporal","Difference","error","(","td","error",")",",","$","\\delta","$",".","We","sample","transition","$","i","$","with","probability",",","$","$","P","(","i",")","=","rac","{","p_i^\u0007lpha","}","{","\\sum_k","p_k^\u0007lpha","}","$","$","where","$","\u0007lpha","$","is","a","hyper-parameter","that","determines","how","much","prioritization","is","used",",","with","$","\u0007lpha","=","0","$","corresponding","to","uniform","case",".","$","p_i","$","is","the","priority",".","We","use","proportional","prioritization","$","p_i","=","|\\delta_i|","+","\\epsilon","$","where","$","\\delta_i","$","is","the","temporal","difference","for","transition","$","i","$",".","We","correct","the","bias","introduced","by","prioritized","replay","using","importance-sampling","(","IS",")","weights","$","$","w_i","=","\bigg","(","rac","{","1","}","{","N","}","rac","{","1","}","{","P","(","i",")","}","\bigg",")","^\beta","$","$","in","the","loss","function",".","This","fully","compensates","when","$","\beta","=","1","$",".","We","normalize","weights","by","$","rac","{","1","}","{","\\max_i","w_i","}","$","for","stability",".","Unbiased","nature","is","most","important","towards","the","convergence","at","end","of","training",".","Therefore","we","increase","$","\beta","$","towards","end","of","training",".","#","#","#","Binary","Segment","Tree","We","use","a","binary","segment","tree","to","efficiently","calculate","$","\\sum_k^i","p_k^\u0007lpha","$",",","the","cumulative","probability",",","which","is","needed","to","sample",".","We","also","use","a","binary","segment","tree","to","find","$","\\min","p_i^\u0007lpha","$",",","which","is","needed","for","$","rac","{","1","}","{","\\max_i","w_i","}","$",".","We","can","also","use","a","min-heap","for","this",".","Binary","Segment","Tree","lets","us","calculate","these","in","$","\\mathcal","{","O","}","(","\\log","n",")","$","time",",","which","is","way","more","efficient","that","the","naive","$","\\mathcal","{","O","}","(","n",")","$","approach",".","This","is","how","a","binary","segment","tree","works","for","sum",";","it","is","similar","for","minimum",".","Let","$","x_i","$","be","the","list","of","$","N","$","values","we","want","to","represent",".","Let","$","b_","{","i",",","j","}","$","be","the","$","j^","{","\\mathop","{","th","}","}","$","node","of","the","$","i^","{","\\mathop","{","th","}","}","$","row","in","the","binary","tree",".","That","is","two","children","of","node","$","b_","{","i",",","j","}","$","are","$","b_","{","i+1,2j","}","$","and","$","b_","{","i+1,2j","+","1","}","$",".","The","leaf","nodes","on","row","$","D","=","\\left\\lceil","{","1","+","\\log_2","N","}","ight","ceil","$","will","have","values","of","$","x","$",".","Every","node","keeps","the","sum","of","the","two","child","nodes",".","That","is",",","the","root","node","keeps","the","sum","of","the","entire","array","of","values",".","The","left","and","right","children","of","the","root","node","keep","the","sum","of","the","first","half","of","the","array","and","the","sum","of","the","second","half","of","the","array",",","respectively",".","And","so","on","...","$","$","b_","{","i",",","j","}","=","\\sum_","{","k","=","(","j","-1",")","*","2^","{","D","-","i","}","+","1","}","^","{","j","*","2^","{","D","-","i","}","}","x_k","$","$","Number","of","nodes","in","row","$","i","$",",","$","$","N_i","=","\\left\\lceil","{","rac","{","N","}","{","D","-","i","+","1","}","}","ight","ceil","$","$","This","is","equal","to","the","sum","of","nodes","in","all","rows","above","$","i","$",".","So","we","can","use","a","single","array","$","a","$","to","store","the","tree",",","where",",","$","$","b_","{","i",",","j","}","ightarrow","a_","{","N_i","+","j","}","$","$","Then","child","nodes","of","$","a_i","$","are","$","a_","{","2i","}","$","and","$","a_","{","2i","+","1","}","$",".","That","is",",","$","$","a_i","=","a_","{","2i","}","+","a_","{","2i","+","1","}","$","$","This","way","of","maintaining","binary","trees","is","very","easy","to","program",".","*","Note","that","we","are","indexing","starting","from","1","*",".","We","use","the","same","structure","to","compute","the","minimum","."]
docstr: ### Initialize
docstr_tok: ["#","#","#","Initialize"]
code: def __init__(self, capacity, alpha):
      """
      ### Initialize
      """
      # We use a power of $2$ for capacity because it simplifies the code and debugging
      self.capacity = capacity
      # $\alpha$
      self.alpha = alpha
  
      # Maintain segment binary trees to take sum and find minimum over a range
      self.priority_sum = [0 for _ in range(2 * self.capacity)]
      self.priority_min = [
          float("inf") for _ in range(2 * self.capacity)
      ]
  
      # Current max priority, $p$, to be assigned to new transitions
      self.max_priority = 1.0
  
      # Arrays for buffer
      self.data = {
          "obs": np.zeros(shape=(capacity, 4, 84, 84), dtype=np.uint8),
          "action": np.zeros(shape=capacity, dtype=np.int32),
          "reward": np.zeros(shape=capacity, dtype=np.float32),
          "next_obs": np.zeros(
              shape=(capacity, 4, 84, 84), dtype=np.uint8
          ),
          "done": np.zeros(shape=capacity, dtype=np.bool),
      }
      # We use cyclic buffers to store data, and `next_idx` keeps the index of the next empty
      # slot
      self.next_idx = 0
  
      # Size of the buffer
      self.size = 0
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/rl/dqn/replay_buffer.py
imports: import random
  import numpy
label_desc: Use image node_id 1 to create a new ReplayBuffer object  with example: obj = ReplayBuffer(capacity, alpha)