Discounted Return of a Trajectory

Problem Statement

Given the sequence of rewards an agent collected during one episode, compute the discounted return from the first time step.

Background

Reinforcement learning agents do not maximise the immediate reward, they maximise the return — the total reward accumulated from now until the end of the episode. Rewards that arrive later are worth less, and the discount factor gamma controls by how much:

$G_t = R_{t+1} + \gamma R_{t+2} + \gamma^2 R_{t+3} + \cdots = \sum_{k=0}^{T-t-1} \gamma^{k} R_{t+k+1}$

Two special cases are worth internalising because they anchor the whole spectrum:

• gamma = 0 makes the agent myopic: the return is exactly the next reward and nothing else matters.
• gamma = 1 makes the agent far-sighted: every reward counts equally, and the return is the plain sum. This is only well defined for episodic tasks that actually terminate.

The recursive form G_t = R_{t+1} + gamma * G_{t+1} is the seed of every algorithm in this plan — the Bellman equations are nothing more than this identity with an expectation wrapped around it.

Your Task

Implement:

def discounted_return(rewards, gamma):
    ...

• rewards — a list of floats, rewards[k] is the reward received at step k (i.e. $R_{k+1}$).
• gamma — the discount factor, 0.0 <= gamma <= 1.0.

Return the discounted return $G_0$ as a float.

Input / Output Format

Input is a Python list of floats and a float. Output is a single float. The grader rounds it to 4 decimal places.

Sample

print(round(discounted_return([1.0, 2.0, 3.0], 0.9), 4))

Output:

5.23

Because 1.0 + 0.92.0 + 0.813.0 = 1.0 + 1.8 + 2.43 = 5.23.