Discounted Return - Reinforcement Learning

Discounted Return

What you'll build

A function that turns a reward sequence into the discounted return G_0, the quantity every RL algorithm tries to maximize.

Background

An agent should prefer reward now over the same reward later, both because the future is uncertain and because an infinite sum of rewards must stay finite. The discount factor γ does this: a reward k steps away is worth γ^k as much as one received now. With γ = 0 the agent is myopic; with γ close to 1 it plans far ahead (typical values: 0.9 to 0.999).

Formula

For rewards r_0, r_1, ..., r_{T-1}:

G_0 = r_0 + γ·r_1 + γ²·r_2 + ... + γ^(T-1)·r_{T-1}
    = Σ_{k=0}^{T-1} γ^k · r_k

Equivalently, backwards: G_t = r_t + γ·G_{t+1} with G_T = 0.

Inputs

• rewards: 1-D array of length T (floats or ints, any sign)
• gamma: float in [0, 1]

Expected return format

A Python float: G_0 rounded to 3 decimals.

Edge cases and conventions

• gamma = 0 returns r_0; gamma = 1 returns the plain sum.
• An empty reward sequence has return 0.0.

Constraints

NumPy only. Vectorize: build the discount vector [γ⁰, γ¹, ..., γ^(T-1)] instead of looping.

Follow-up

• Return the whole vector [G_0, ..., G_{T-1}] in O(T) (hint: the backward recursion). REINFORCE needs exactly this.
• Why does γ < 1 guarantee a finite return for bounded rewards in an infinite-horizon task?

Input: rewards = [1, 1, 1, 1], gamma = 0.9
Output: 3.439

Input: rewards = [0, 0, 0, 10], gamma = 0.5
Output: 1.25

import numpy as np

def discounted_return(rewards: np.ndarray, gamma: float) -> float:
    """
    Calculate the discounted return G_0.

    Args:
        rewards: (T,) rewards [r_0, r_1, ..., r_{T-1}]
        gamma: discount factor in [0, 1]

    Returns:
        G_0 = sum_k gamma**k * r_k as a float, rounded to 3 decimals
    """
    rewards = np.array(rewards)
    # Your code here
    pass