Discounted Return of a Trajectory

easyFREEReturns, Policies, and Value FunctionsReinforcement LearningRL 1: Foundations

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.

Example

Input

discounted_return([1.0, 2.0, 3.0], 0.9)

Output

5.23

Reasoning

The discounted return $G_0$ is calculated by summing each reward weighted by the discount factor $\gamma$ raised to the power of its time step index: $G_0 = \sum_{k=0}^{n-1} \gamma^k R_k$ .
For the first reward $R_0 = 1.0$ at time step $k=0$ , the contribution is $0.9^0 \times 1.0 = 1 \times 1.0 = 1.0$ .
For the second reward $R_1 = 2.0$ at time step $k=1$ , the contribution is $0.9^1 \times 2.0 = 0.9 \times 2.0 = 1.8$ .
For the third reward $R_2 = 3.0$ at time step $k=2$ , the contribution is $0.9^2 \times 3.0 = 0.81 \times 3.0 = 2.43$ .
Summing these weighted contributions gives the total return: $1.0 + 1.8 + 2.43 = 5.23$ .
The final output is 5.23

Constraints

0 <= len(rewards) <= 1000
0.0 <= gamma <= 1.0
Rewards may be negative.
An empty reward list has return 0.0.
Return a plain Python float; do not round inside the function.

Learn this in

Introduction to MLML Study Plan Reinforcement LearningML Study Plan Introduction to Reinforcement LearningReinforcement Learning Markov Decision ProcessesReinforcement Learning

More from RL 1: Foundations

GAE Over a Full Trajectoryhard PPO Clipped Objective (Single Sample)medium TD Residuals for GAEmedium Return Standardization Baselinemedium REINFORCE Loss Termmedium

Concept primer

1. Background Knowledge

In Reinforcement Learning (RL), an agent interacts with an environment over time to maximize a cumulative signal called the return. Unlike supervised learning, where targets are static, RL deals with sequential decision-making where current actions affect future states and rewards. The core metric for success is not the immediate reward $R_{t+1}$ , but the discounted return $G_t$ , which sums all future rewards weighted by a discount factor $\gamma$ (gamma).

The discount factor $\gamma \in [0, 1]$ determines the present value of future rewards. A $\gamma$ close to 0 makes the agent myopic, caring only about immediate gains. A $\gamma$ close to 1 makes the agent far-sighted, valuing long-term consequences. Mathematically, the return from time step $t$ is defined as:

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

This formulation ensures that rewards further in the future contribute less to the current decision, promoting stability in learning algorithms and ensuring convergence in infinite-horizon tasks. Understanding this weighted sum is fundamental to deriving Bellman equations, which form the basis of value-based RL methods like Q-Learning and SARSA.

2. Algorithm Approach

The problem requires computing a weighted sum of a sequence. While the definition involves powers of $\gamma$ ( $\gamma^0, \gamma^1, \gamma^2, \dots$ ), calculating these powers explicitly using pow(gamma, k) or gamma ** k is computationally inefficient and prone to floating-point precision errors for long sequences.

A more robust approach leverages the recursive structure of the return. Notice that:

G_t = R_{t+1} + \gamma (R_{t+2} + \gamma R_{t+3} + \dots) = R_{t+1} + \gamma G_{t+1}

This implies we can compute the return backwards from the last time step to the first. Starting from the end of the episode, the return is simply the last reward. We then iteratively update the return by adding the current reward and multiplying the accumulated future return by $\gamma$ . This backward induction method avoids repeated exponentiation and reduces the operation count significantly.

3. Step-by-Step Strategy

Initialize Accumulator: Create a variable, say G, initialized to 0.0. This will hold the discounted return from the current step to the end.
Iterate Backwards: Loop through the rewards list in reverse order (from the last element to the first). This is crucial because the return at step $t$ depends on the return at step $t+1$ .
Update Return: For each reward $R$ encountered in the reverse loop:

Multiply the current accumulated return G by gamma.
Add the current reward $R$ to this product.
Update G with this new value: G = R + gamma * G.

Return Result: After processing all rewards, G will contain $G_0$ , the discounted return from the first time step. Return this value.

Example Trace: For rewards = [1.0, 2.0, 3.0] and gamma = 0.9:

Start: G = 0.0
Step 1 (Reward 3.0): G = 3.0 + 0.9 * 0.0 = 3.0
Step 2 (Reward 2.0): G = 2.0 + 0.9 * 3.0 = 4.7
Step 3 (Reward 1.0): G = 1.0 + 0.9 * 4.7 = 5.23

4. Common Pitfalls

Forward Iteration with Powers: Attempting to iterate forward and calculating gamma ** k for each term is inefficient ( $O(N)$ power calculations) and less numerically stable. The backward pass is $O(1)$ per step.
Off-by-One Errors: Ensure you iterate through all rewards. The return $G_0$ includes $R_1$ (the first reward in the list). Missing the first or last reward will yield an incorrect result.
Gamma Edge Cases:
If gamma = 0, the return should be just the first reward. The backward algorithm handles this naturally: G becomes R_last, then R_prev + 0, etc., eventually leaving only R_first.
If gamma = 1, the return is the simple sum. The algorithm reduces to G = R + G, which is a standard sum.
Empty Input: Although the problem implies a valid episode, consider what happens if rewards is empty. The loop won't execute, and G remains 0.0, which is mathematically correct for an empty sum.
Floating Point Precision: While Python floats are double-precision, accumulating many small errors can occur. The backward method is generally more stable than forward exponentiation.

5. Time & Space Complexity

Time Complexity: $O(N)$ , where $N$ is the length of the rewards list. We iterate through the list exactly once, performing constant-time arithmetic operations (multiplication and addition) at each step.
Space Complexity: $O(1)$ auxiliary space. We only need a single variable to store the accumulated return G. No additional data structures proportional to the input size are required.

solution.py

Python 3.13

Test results