Ước Lượng Lợi Thế Tổng Quát (Generalized Advantage Estimation - GAE)

Mục tiêu kiến thức

• Trong RL, ta cần ước lượng lợi thế A(s_t, a_t) = Q(s_t, a_t) - V(s_t).
• Phương pháp 1 bước TD error: δ_t = r_t + γ V(s_t+1) - V(s_t) (bias cao do dựa vào V, phương sai thấp).
• Phương pháp Monte Carlo: G_t - V(s_t) (không có bias, nhưng phương sai cực lớn).
• Schulman (2015) đề xuất GAE(γ, λ) sử dụng trung bình có trọng số hàm mũ của các bước TD:

δ_t^V = r_t + γ V(s_t+1)(1 - d_t) - V(s_t) A_t^GAE = ∑_l=0^∞ (γ λ)^l δ_t+l^V = δ_t^V + γ λ (1 - d_t) A_t+1^GAE Trong đó:

• d_t ∈ {0, 1} là cờ kết thúc episode (done). Nếu d_t = 1, bước tiếp theo bị ngắt kết nối.
• λ = 0: Trở thành sai số TD(0) 1 bước.
• λ = 1: Trở thành ước lượng Monte Carlo đầy đủ.
• λ ∈ [0.9, 0.95] mang lại sự cân bằng hoàn hảo giữa bias và variance, là chuẩn mực trong PPO.

Yêu cầu

Viết hàm:

def compute_gae(
    rewards: torch.Tensor,
    values: torch.Tensor,
    next_value: torch.Tensor,
    dones: torch.Tensor,
    gamma: float = 0.99,
    lam: float = 0.95
) -> tuple[torch.Tensor, torch.Tensor]:
    pass

• rewards: Tensor 1D (T,).
• values: Tensor 1D (T,) chứa ước lượng V(s_t).
• next_value: Scalar tensor giá trị V(s_T).
• dones: Tensor 1D (T,) kiểu boolean hoặc float (1.0 nếu kết thúc, 0.0 nếu tiếp tục).
• Trả về tuple (advantages, returns):
• advantages: Tensor 1D (T,) chứa A_t^GAE.
• returns: Tensor 1D (T,) được tính là A_t + V(s_t) (dùng làm target huấn luyện mạng Critic).

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.