Discounted Returns

Compute the discounted return at every step of an episode.

An agent collects a reward r_t at each step. The return G_t from step t is the sum of all future rewards, with each step further away multiplied by one more factor of gamma:

G_t = r_t + gamma * r_{t+1} + gamma^2 * r_{t+2} + ...

Computing each G_t from scratch costs O(T^2). Work backwards instead, with G_t = r_t + gamma * G_{t+1} and G_T = 0. Policy gradient methods such as REINFORCE weight each action by exactly this quantity.

Function signature

def discounted_returns(rewards, gamma):

Example

Input

discounted_returns(np.array([1.0, 0.0, 0.0, 1.0]), gamma=0.9)

Output

array([1.729, 0.81 , 0.9  , 1.   ])

Reasoning

Start at the end. G_3 = 1. G_2 = 0 + 0.9 * 1 = 0.9. G_1 = 0 + 0.9 * 0.9 = 0.81. G_0 = 1 + 0.9 * 0.81 = 1.729.

Constraints and notes

• rewards is a 1D array. It can be empty.
• 0 <= gamma <= 1.
• Return a float array with the same length as rewards.