Return Variance Under Stochastic Rewards

hardFREEReturns, Policies, and Value FunctionsReinforcement LearningRL 1: Foundations

Problem Statement

Each timestep t gives an independent reward with mean mu[t] and variance var[t]. The discounted return is G = sum_t gamma^t * R_t. Because the rewards are independent, the variance of the discounted return is:

$\mathrm{Var}(G) = \sum_{t=0}^{T-1} \gamma^{2t}\, \mathrm{Var}(R_t)$

Implement return_variance(var, gamma) returning Var(G) as a float. (mu is not needed for the variance.)

Example

Input

return_variance([1.0, 1.0], 0.5)

Output

1.25

Reasoning

Initialize the cumulative variance to $0.0$ and the discount factor for the current step to $1.0$ . Since the variance formula requires squaring the discount rate at each step, calculate the squared discount factor: $g^2 = 0.5^2 = 0.25$ .
Process the first variance value ( $1.0$ ) at timestep $t=0$ . Multiply the current discount factor by this variance and add it to the total: $0.0 + (1.0 \times 1.0) = 1.0$ .
Update the discount factor for the next timestep by multiplying it by the squared discount factor: $1.0 \times 0.25 = 0.25$ .
Process the second variance value ( $1.0$ ) at timestep $t=1$ . Multiply the updated discount factor by this variance and add it to the running total: $1.0 + (0.25 \times 1.0) = 1.25$ .
The final output is 1.25

Constraints

1 <= len(var) <= 10000, all var[t] >= 0, 0.0 <= gamma <= 1.0
Note the discount is squared: gamma**(2t).
Return a float.

Learn this in

Monte Carlo & Temporal-Difference LearningReinforcement 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, the discounted return $G$ is defined as the sum of future rewards, each scaled by a discount factor $\gamma^t$ where $0 \le \gamma < 1$ . This discounting ensures that $G$ remains finite even in infinite-horizon tasks and reflects the principle that immediate rewards are more valuable than distant ones. The return is a random variable because the rewards $R_t$ are stochastic.

A fundamental property of probability theory states that for independent random variables, the variance of their sum equals the sum of their variances. Specifically, if $X$ and $Y$ are independent, then $\text{Var}(aX + bY) = a^2\text{Var}(X) + b^2\text{Var}(Y)$ . Note that the scaling factor is squared. This is distinct from the mean, where $\text{E}[aX + bY] = a\text{E}[X] + b\text{E}[Y]$ . The mean $\mu$ does not appear in the variance calculation because variance measures spread around the mean, not the location of the mean itself.

Therefore, for the discounted return $G = \sum_{t=0}^{T-1} \gamma^t R_t$ with independent rewards, the variance decomposes as:

\text{Var}(G) = \sum_{t=0}^{T-1} \gamma^{2t} \text{Var}(R_t)

This formula is central to understanding risk in RL: even if the expected return is high, a high variance indicates unpredictable outcomes, which is critical for risk-sensitive policy optimization.

2. Algorithm Approach

The problem reduces to a straightforward weighted sum computation. You are given an array of variances var and a scalar gamma. The goal is to compute:

\sum_{t=0}^{T-1} (\gamma^2)^t \cdot \text{var}[t]

This is a geometric-series-like weighted sum where the weight for each term is $\gamma^{2t}$ . The most direct approach is to iterate through the var array, maintaining a running power of $\gamma^2$ , and accumulating the product with each variance value. No complex data structures or dynamic programming are needed; this is a linear scan with constant extra state.

3. Step-by-Step Strategy

Initialize a variable total_variance to 0.0 and a variable power to 1.0. The power variable will track $\gamma^{2t}$ .
Iterate over each element v in the var array with index t.
Accumulate: Add power * v to total_variance.
Update power: Multiply power by $\gamma^2$ for the next iteration. This avoids recalculating $\gamma^{2t}$ from scratch each time, which would be less efficient and more prone to floating-point drift.
Return total_variance as a float.

Pseudocode:

def return_variance(var, gamma):
 total = 0.0
 power = 1.0
 gamma_sq = gamma * gamma
 for v in var:
 total += power * v
 power *= gamma_sq
 return total

4. Common Pitfalls

Forgetting to square gamma: The variance formula uses $\gamma^{2t}$ , not $\gamma^t$ . Using $\gamma^t$ will give an incorrect result. Always square the discount factor when dealing with variance.
Using mu unnecessarily: The problem explicitly states that mu is not needed. Including it in the calculation will lead to errors. Variance is independent of the mean.
Floating-point precision: For large $T$ and $\gamma$ close to 1, $\gamma^{2t}$ may underflow to 0.0 in floating-point arithmetic. This is generally acceptable for practical purposes, but be aware that very small terms may be lost.
Off-by-one errors: Ensure you iterate over all elements in var. The sum is from $t=0$ to $T-1$ , where $T = \text{len(var)}$ .
Negative gamma: While $\gamma$ is typically in $[0, 1)$ , if negative values are allowed, $\gamma^2$ is still positive, so the formula holds. However, standard RL assumes $0 \le \gamma < 1$ .

5. Time & Space Complexity

Time Complexity: $O(T)$ , where $T$ is the length of the var array. Each element is visited exactly once, and each iteration performs a constant number of arithmetic operations.
Space Complexity: $O(1)$ extra space. Only a few scalar variables (total, power, gamma_sq) are used, regardless of the input size. The input array is not modified.

solution.py

Python 3.13

Test results