Problems
Loading...
1 / 1

Discounted Returns

Reinforcement Learning
Easy

In reinforcement learning, an agent collects rewards at each timestep. The discounted return at timestep t is the sum of all future rewards, where each reward is discounted by a factor of gamma raised to the power of how far in the future it occurs. This captures the idea that immediate rewards are worth more than distant ones.

Given a list of rewards collected over T timesteps and a discount factor gamma, compute the discounted return for every timestep.

Formula

The discounted return at timestep t is:

Gt=rt+γrt+1+γ2rt+2++γT1trT1G_t = r_t + \gamma \cdot r_{t+1} + \gamma^2 \cdot r_{t+2} + \cdots + \gamma^{T-1-t} \cdot r_{T-1}

This can be computed efficiently using the backward recursive relation:

Gt=rt+γGt+1G_t = r_t + \gamma \cdot G_{t+1} GT1=rT1G_{T-1} = r_{T-1}
Loading visualization...

Examples

Input:

rewards = [1, 1, 1], gamma = 1.0

Output:

[3.0, 2.0, 1.0]

With gamma = 1 there is no discounting. G[0] = 1 + 1 + 1 = 3, G[1] = 1 + 1 = 2, G[2] = 1.

Input:

rewards = [0, 0, 0, 10], gamma = 0.9

Output:

[7.29, 8.1, 9.0, 10.0]

Only the last timestep has a reward. The return at earlier steps is that reward discounted by gamma raised to the number of steps away: G[0] = 10 * 0.9^3 = 7.29.

Hint 1

Initialize an array G of the same length as rewards. Set G[T-1] = rewards[T-1]. Then loop from t = T-2 down to 0, computing G[t] = rewards[t] + gamma * G[t+1].

Hint 2

The key insight is to process the rewards from right to left. Each return only depends on the return one step ahead, which you have already computed.

Requirements

  • Compute the discounted return for every timestep using backward recursion
  • Start from the last timestep where G[T-1] = rewards[T-1]
  • For each earlier timestep: G[t] = rewards[t] + gamma * G[t+1]
  • Return a list of floats with the same length as rewards

Constraints

  • rewards has at least one element
  • 0 <= gamma <= 1
  • Return a list of floats with the same length as rewards
  • Time limit: 300 ms
Try Similar Problems