class_docstr: An N-step transition adder.
  
  This will buffer a sequence of N timesteps in order to form a single N-step
  transition which is added to reverb for future retrieval.
  
  For N=1 the data added to replay will be a standard one-step transition which
  takes the form:
  
        (s_t, a_t, r_t, d_t, s_{t+1}, e_t)
  
  where:
  
    s_t = state observation at time t
    a_t = the action taken from s_t
    r_t = reward ensuing from action a_t
    d_t = environment discount ensuing from action a_t. This discount is
        applied to future rewards after r_t.
    e_t [Optional] = extra data that the agent persists in replay.
  
  For N greater than 1, transitions are of the form:
  
        (s_t, a_t, R_{t:t+n}, D_{t:t+n}, s_{t+N}, e_t),
  
  where:
  
    s_t = State (observation) at time t.
    a_t = Action taken from state s_t.
    g = the additional discount, used by the agent to discount future returns.
    R_{t:t+n} = N-step discounted return, i.e. accumulated over N rewards:
          R_{t:t+n} := r_t + g * d_t * r_{t+1} + ...
                           + g^{n-1} * d_t * ... * d_{t+n-2} * r_{t+n-1}.
    D_{t:t+n}: N-step product of agent discounts g_i and environment
      "discounts" d_i.
          D_{t:t+n} := g^{n-1} * d_{t} * ... * d_{t+n-1},
      For most environments d_i is 1 for all steps except the last,
      i.e. it is the episode termination signal.
    s_{t+n}: The "arrival" state, i.e. the state at time t+n.
    e_t [Optional]: A nested structure of any 'extras' the user wishes to add.
  
  Notes:
    - At the beginning and end of episodes, shorter transitions are added.
      That is, at the beginning of the episode, it will add:
            (s_0 -> s_1), (s_0 -> s_2), ..., (s_0 -> s_n), (s_1 -> s_{n+1})
  
      And at the end of the episode, it will add:
            (s_{T-n+1} -> s_T), (s_{T-n+2} -> s_T), ... (s_{T-1} -> s_T).
    - We add the *first* `extra` of each transition, not the *last*, i.e.
        if extras are provided, we get e_t, not e_{t+n}.
class_docstr_tok: ["An","N-step","transition","adder",".","This","will","buffer","a","sequence","of","N","timesteps","in","order","to","form","a","single","N-step","transition","which","is","added","to","reverb","for","future","retrieval",".","For","N=1","the","data","added","to","replay","will","be","a","standard","one-step","transition","which","takes","the","form",":","(","s_t",",","a_t",",","r_t",",","d_t",",","s_","{","t+1","}",",","e_t",")","where",":","s_t","=","state","observation","at","time","t","a_t","=","the","action","taken","from","s_t","r_t","=","reward","ensuing","from","action","a_t","d_t","=","environment","discount","ensuing","from","action","a_t",".","This","discount","is","applied","to","future","rewards","after","r_t",".","e_t","[","Optional","]","=","extra","data","that","the","agent","persists","in","replay",".","For","N","greater","than","1",",","transitions","are","of","the","form",":","(","s_t",",","a_t",",","R_","{","t",":","t+n","}",",","D_","{","t",":","t+n","}",",","s_","{","t+N","}",",","e_t",")",",","where",":","s_t","=","State","(","observation",")","at","time","t.","a_t","=","Action","taken","from","state","s_t",".","g","=","the","additional","discount",",","used","by","the","agent","to","discount","future","returns",".","R_","{","t",":","t+n","}","=","N-step","discounted","return",",","i.e",".","accumulated","over","N","rewards",":","R_","{","t",":","t+n","}",":","=","r_t","+","g","*","d_t","*","r_","{","t+1","}","+","...","+","g^","{","n-1","}","*","d_t","*","...","*","d_","{","t+n-2","}","*","r_","{","t+n-1","}",".","D_","{","t",":","t+n","}",":","N-step","product","of","agent","discounts","g_i","and","environment","``","discounts","''","d_i",".","D_","{","t",":","t+n","}",":","=","g^","{","n-1","}","*","d_","{","t","}","*","...","*","d_","{","t+n-1","}",",","For","most","environments","d_i","is","1","for","all","steps","except","the","last",",","i.e",".","it","is","the","episode","termination","signal",".","s_","{","t+n","}",":","The","``","arrival","''","state",",","i.e",".","the","state","at","time","t+n",".","e_t","[","Optional","]",":","A","nested","structure","of","any","'extras","'","the","user","wishes","to","add",".","Notes",":","-","At","the","beginning","and","end","of","episodes",",","shorter","transitions","are","added",".","That","is",",","at","the","beginning","of","the","episode",",","it","will","add",":","(","s_0","-",">","s_1",")",",","(","s_0","-",">","s_2",")",",","...",",","(","s_0","-",">","s_n",")",",","(","s_1","-",">","s_","{","n+1","}",")","And","at","the","end","of","the","episode",",","it","will","add",":","(","s_","{","T-n+1","}","-",">","s_T",")",",","(","s_","{","T-n+2","}","-",">","s_T",")",",","...","(","s_","{","T-1","}","-",">","s_T",")",".","-","We","add","the","*","first","*","`","extra","`","of","each","transition",",","not","the","*","last","*",",","i.e",".","if","extras","are","provided",",","we","get","e_t",",","not","e_","{","t+n","}","."]
docstr: Creates an N-step transition adder.
  
  Args:
    client: A `reverb.Client` to send the data to replay through.
    n_step: The "N" in N-step transition. See the class docstring for the
      precise definition of what an N-step transition is. `n_step` must be at
      least 1, in which case we use the standard one-step transition, i.e.
      (s_t, a_t, r_t, d_t, s_t+1, e_t).
    discount: Discount factor to apply. This corresponds to the agent's
      discount in the class docstring.
    priority_fns: See docstring for BaseAdder.
    max_in_flight_items: The maximum number of items allowed to be "in flight"
      at the same time. See `block_until_num_items` in
      `reverb.TrajectoryWriter.flush` for more info.
  
  Raises:
    ValueError: If n_step is less than 1.
docstr_tok: ["Creates","an","N-step","transition","adder",".","Args",":","client",":","A","`","reverb.Client","`","to","send","the","data","to","replay","through",".","n_step",":","The","``","N","''","in","N-step","transition",".","See","the","class","docstring","for","the","precise","definition","of","what","an","N-step","transition","is",".","`","n_step","`","must","be","at","least","1",",","in","which","case","we","use","the","standard","one-step","transition",",","i.e",".","(","s_t",",","a_t",",","r_t",",","d_t",",","s_t+1",",","e_t",")",".","discount",":","Discount","factor","to","apply",".","This","corresponds","to","the","agent's","discount","in","the","class","docstring",".","priority_fns",":","See","docstring","for","BaseAdder",".","max_in_flight_items",":","The","maximum","number","of","items","allowed","to","be","``","in","flight","''","at","the","same","time",".","See","`","block_until_num_items","`","in","`","reverb.TrajectoryWriter.flush","`","for","more","info",".","Raises",":","ValueError",":","If","n_step","is","less","than","1","."]
code: def __init__(
      self,
      client: reverb.Client,
      n_step: int,
      discount: float,
      *,
      priority_fns: Optional[base.PriorityFnMapping] = None,
      max_in_flight_items: int = 5,
  ):
      """Creates an N-step transition adder.
  
      Args:
        client: A `reverb.Client` to send the data to replay through.
        n_step: The "N" in N-step transition. See the class docstring for the
          precise definition of what an N-step transition is. `n_step` must be at
          least 1, in which case we use the standard one-step transition, i.e.
          (s_t, a_t, r_t, d_t, s_t+1, e_t).
        discount: Discount factor to apply. This corresponds to the agent's
          discount in the class docstring.
        priority_fns: See docstring for BaseAdder.
        max_in_flight_items: The maximum number of items allowed to be "in flight"
          at the same time. See `block_until_num_items` in
          `reverb.TrajectoryWriter.flush` for more info.
  
      Raises:
        ValueError: If n_step is less than 1.
      """
      # Makes the additional discount a float32, which means that it will be
      # upcast if rewards/discounts are float64 and left alone otherwise.
      self.n_step = n_step
      self._discount = tree.map_structure(np.float32, discount)
      self._first_idx = 0
      self._last_idx = 0
  
      super().__init__(
          client=client,
          max_sequence_length=n_step + 1,
          priority_fns=priority_fns,
          max_in_flight_items=max_in_flight_items,
      )
  
imports: import copy
  from typing import Optional, Tuple
  from acme import specs
  from acme import types
  from acme.adders.reverb import base
  from acme.adders.reverb import utils
  from acme.utils import tree_utils
  import numpy
  import reverb
  import tree
label_desc: Use image node_id 1 to create a new NStepTransitionAdder object from inherited base classes: base with example: obj = NStepTransitionAdder(client, n_step, discount)