parameters: (reward, value, length, discount, window)
return_statement: return tf.check_numerics(tf.stop_gradient(mask * return_), 'return')
docstring: N-step discounted return.
docstring_summary: N-step discounted return.
docstring_tokens: N - step discounted return .
function: def fixed_step_return(reward, value, length, discount, window):
    """N-step discounted return."""
    timestep = tf.range(reward.shape[1].value)
    mask = tf.cast(timestep[None, :] < length[:, None], tf.float32)
    return_ = tf.zeros_like(reward)
    for _ in range(window):
      return_ += reward
      reward = discount * tf.concat([reward[:, 1:], tf.zeros_like(reward[:, -1:])], 1)
    return_ += discount**window * tf.concat(
        [value[:, window:], tf.zeros_like(value[:, -window:]), 1])
    return tf.check_numerics(tf.stop_gradient(mask * return_), 'return')
function_tokens: def fixed_step_return ( reward , value , length , discount , window ) : timestep = tf . range ( reward . shape [ 1 ] . value ) mask = tf . cast ( timestep [ None , : ] < length [ : , None ] , tf . float32 ) return_ = tf . zeros_like ( reward ) for _ in range ( window ) : return_ += reward reward = discount * tf . concat ( [ reward [ : , 1 : ] , tf . zeros_like ( reward [ : , - 1 : ] ) ] , 1 ) return_ += discount ** window * tf . concat ( [ value [ : , window : ] , tf . zeros_like ( value [ : , - window : ] ) , 1 ] ) return tf . check_numerics ( tf . stop_gradient ( mask * return_ ) , 'return' )