import tensorflow as tf

    reward = np.reshape(reward, ((batch_size, trace_length)))
    reward_buffer = np.zeros(((batch_size, trace_length+1)))
    reward_buffer[:, :trace_length] = reward
    discounted_reward = np.zeros(((batch_size, trace_length)))

    for t in range(trace_length-1, -1, -1):
        reward_buffer[:,t+1:] *= y
        discounted_reward[:,t] = np.sum(reward_buffer[:,t:],1)

    return np.reshape(discounted_reward,(batch_size *trace_length))


def make_cube(trace_length):
    cube = tf.Variable(tf.zeros([trace_length, trace_length, trace_length]))

    cube_ops = []
    for i in range(trace_length):
        cube_ops.append(cube[i, :(i+1), :(i+1)].assign(tf.ones([i+1, i+1])))
    return cube, cube_ops
