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","}","."]
code: def _compute_cumulative_quantities(
      self, rewards: types.NestedArray, discounts: types.NestedArray
  ) -> Tuple[types.NestedArray, types.NestedArray]:
      # Give the same tree structure to the n-step return accumulator,
      # n-step discount accumulator, and self.discount, so that they can be
      # iterated in parallel using tree.map_structure.
      (
          rewards,
          discounts,
          self_discount,
      ) = tree_utils.broadcast_structures(
          rewards, discounts, self._discount
      )
      flat_rewards = tree.flatten(rewards)
      flat_discounts = tree.flatten(discounts)
      flat_self_discount = tree.flatten(self_discount)
  
      # Copy total_discount as it is otherwise read-only.
      total_discount = [np.copy(a[0]) for a in flat_discounts]
  
      # Broadcast n_step_return to have the broadcasted shape of
      # reward * discount.
      n_step_return = [
          np.copy(np.broadcast_to(r[0], np.broadcast(r[0], d).shape))
          for r, d in zip(flat_rewards, total_discount)
      ]
  
      # NOTE: total_discount will have one less self_discount applied to it than
      # the value of self._n_step. This is so that when the learner/update uses
      # an additional discount we don't apply it twice. Inside the following loop
      # we will apply this right before summing up the n_step_return.
      for i in range(1, self._n_step):
          for nsr, td, r, d, sd in zip(
              n_step_return,
              total_discount,
              flat_rewards,
              flat_discounts,
              flat_self_discount,
          ):
              # Equivalent to: `total_discount *= self._discount`.
              td *= sd
              # Equivalent to: `n_step_return += reward[i] * total_discount`.
              nsr += r[i] * td
              # Equivalent to: `total_discount *= discount[i]`.
              td *= d[i]
  
      n_step_return = tree.unflatten_as(rewards, n_step_return)
      total_discount = tree.unflatten_as(rewards, total_discount)
      return n_step_return, total_discount
  
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 7 for calling the NStepTransitionAdder obj's underlying member method code with example usage: obj._compute_cumulative_quantities(rewards, discounts) and returns: n_step_return,  total_discount