code: def _n_step_transition_from_episode(
      observations: acme_types.NestedTensor,
      actions: tf.Tensor,
      rewards: tf.Tensor,
      discounts: tf.Tensor,
      n_step: int,
      discount: float,
  ):
      """Produce Reverb-like N-step transition from a full episode.
  
      Observations, actions, rewards and discounts have the same length. This
      function will ignore the first reward and discount and the last action.
  
      Args:
        observations: [L, ...] Tensor.
        actions: [L, ...] Tensor.
        rewards: [L] Tensor.
        discounts: [L] Tensor.
        n_step: number of steps to squash into a single transition.
        discount: discount to use for TD updates.
  
      Returns:
        (o_t, a_t, r_t, d_t, o_tp1) tuple.
      """
  
      max_index = tf.shape(rewards)[0] - 1
      first = tf.random.uniform(
          shape=(), minval=0, maxval=max_index - 1, dtype=tf.int32
      )
      last = tf.minimum(first + n_step, max_index)
  
      o_t = tree.map_structure(operator.itemgetter(first), observations)
      a_t = tree.map_structure(operator.itemgetter(first), actions)
      o_tp1 = tree.map_structure(
          operator.itemgetter(last), observations
      )
  
      # 0, 1, ..., n-1.
      discount_range = tf.cast(tf.range(last - first), tf.float32)
      # 1, g, ..., g^{n-1}.
      additional_discounts = tf.pow(discount, discount_range)
      # 1, d_t, d_t * d_{t+1}, ..., d_t * ... * d_{t+n-2}.
      discounts = tf.concat(
          [[1.0], tf.math.cumprod(discounts[first : last - 1])], 0
      )
      # 1, g * d_t, ..., g^{n-1} * d_t * ... * d_{t+n-2}.
      discounts *= additional_discounts
      # r_t + g * d_t * r_{t+1} + ... + g^{n-1} * d_t * ... * d_{t+n-2} * r_{t+n-1}
      # We have to shift rewards by one so last=max_index corresponds to transitions
      # that include the last reward.
      r_t = tf.reduce_sum(rewards[first + 1 : last + 1] * discounts)
  
      # g^{n-1} * d_{t} * ... * d_{t+n-1}.
      d_t = discounts[-1]
  
      info = tree.map_structure(
          lambda dtype: tf.ones([], dtype),
          reverb.SampleInfo.tf_dtypes(),
      )
      return reverb.ReplaySample(
          info=info,
          data=acme_types.Transition(o_t, a_t, r_t, d_t, o_tp1),
      )
  
imports: import copy
  import functools
  import operator
  from typing import Optional
  from acme import datasets
  from acme import specs
  from acme import types
  from acme.adders import reverb
  from acme.agents import agent
  from acme.agents.tf import actors
  from acme.agents.tf import dqn
  from acme.tf import utils
  import reverb
  import sonnet
  import tensorflow
  import tree
  import trfl
label_desc: Use image node_id 1 for calling a global function with example usage: _n_step_transition_from_episode(observations, actions, rewards, discounts, n_step, discount) and returns: reverb