RL中高效累计回报（Return）计算Python代码问题求助

问题描述 #

以下用于计算强化学习累计折扣回报的代码运行后结果要么为0要么是空列表，运行测试代码时触发断言错误：

原代码：

import numpy as np
def discounted_return(rewards: np.ndarray, gamma: float) -> np.ndarray:
    '''
    Computes all returns for the given sequence of rewards.
    :param rewards: The sequence of rewards as a np.ndarray.
    :param gamma: The discount factor as a `float`.
    :returns: The discounted return for each time step; as a `np.ndarray`.
    ''' 
    T = len(rewards)
    
    returns = np.zeros_like(T, dtype=float)
    
    # implement the efficient return computation
    # YOUR CODE HERE
    
    R = 0
    for t in reversed(range(1, T-1)):
        # update the total discounted reward
        R = R * gamma + rewards[t]
        returns[t] = R
    
    return returns

测试代码：

print(discounted_return(np.array([0, 1]), 0.9))
assert np.isclose(np.array([1.0, 0]), discounted_return(np.array([0, 1]), 0.9)).all()# , "`discounted_return` is not implemented or wrong."

错误点分析 #

• 数组初始化错误：np.zeros_like(T)中T是整数，zeros_like会生成一个标量而非数组，导致后续赋值操作完全无效。正确方式是创建长度为T的全0数组，比如np.zeros(T, dtype=float)或np.zeros_like(rewards)。
• 循环范围错误：reversed(range(1, T-1))的范围完全错误。当T=2时，range(1,1)是空序列，循环根本不会执行；即使T更大，也会漏掉首尾的时间步。正确逻辑是覆盖所有时间步，从最后一个索引（T-1）倒序遍历到0。
• 回报计算逻辑遗漏：最后一个时间步的回报就是其本身的奖励，循环需要从该位置开始向前递推计算。

修复后的代码 #

import numpy as np
def discounted_return(rewards: np.ndarray, gamma: float) -> np.ndarray:
    '''
    Computes all returns for the given sequence of rewards.
    :param rewards: The sequence of rewards as a np.ndarray.
    :param gamma: The discount factor as a `float`.
    :returns: The discounted return for each time step; as a `np.ndarray`.
    ''' 
    T = len(rewards)
    # 正确初始化长度为T的全0数组
    returns = np.zeros(T, dtype=float)
    
    R = 0.0
    # 从最后一个时间步倒序遍历到第0步
    for t in range(T-1, -1, -1):
        R = rewards[t] + gamma * R
        returns[t] = R
    
    return returns

测试验证 #

运行测试代码时需注意：原断言的预期值np.array([1.0, 0])是错误的，正确的累计折扣回报应为：

• t=0时：$G_0 = r_0 + \gamma r_1 = 0 + 0.9*1 = 0.9$
• t=1时：$G_1 = r_1 = 1.0$

修正后的测试代码：

print(discounted_return(np.array([0, 1]), 0.9))  # 输出 [0.9 1. ]
assert np.isclose(np.array([0.9, 1.0]), discounted_return(np.array([0, 1]), 0.9)).all()

内容的提问来源于stack exchange，提问作者Gunners