Thuật Toán Gradient Chính Sách REINFORCE Kèm Baseline Giảm Phương Sai

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

• Định lý Gradient Chính sách (Policy Gradient Theorem - Williams 1992):

∇_theta J(theta) = E_tau [ Σ_t=0^T ∇_theta log pi_theta(a_t | s_t) (G_t - b(s_t)) ] Trong đó:

• G_t = Σ_k=t^T gamma^k-t r_k là tổng phần thưởng chiết khấu (Return) từ bước t.
• b(s_t) là hàm đường cơ sở (baseline, ví dụ ước lượng giá trị trạng thái V(s_t)).
• Trừ b(s_t) giảm đáng kể phương sai (variance reduction) của ước lượng Monte Carlo mà không làm thay đổi kỳ vọng gradient vì:

E_{at sim pi} [∇_theta log pi_theta(a_t | s_t) b(s_t)] = b(s_t) ∇_theta Σ_a pi_theta(a | s_t) = b(s_t) ∇_theta 1 = 0

• Vì các optimizer (Adam, SGD) trong PyTorch thực hiện cực tiểu hóa (loss minimization), hàm mục tiêu mất mát REINFORCE là:

L_policy = -Σ_t=0^T log pi_theta(a_t | s_t) · (G_t - b_t)

Yêu cầu

Viết hàm:

def compute_reinforce_loss(
    log_probs: torch.Tensor,
    returns: torch.Tensor,
    baseline_values: torch.Tensor | None = None
) -> torch.Tensor:
    pass

• log_probs: Tensor 1D (T,) chứa log pi(a_t | s_t).
• returns: Tensor 1D (T,) chứa G_t.
• baseline_values: Tensor 1D (T,) chứa b_t, hoặc None nếu không có baseline (khi đó b_t = 0).
• Quan trọng: baseline_values phải được tách khỏi đồ thị gradient (detach()) để gradient của policy loss không lan truyền ngược làm hỏng baseline.
• Trả về scalar tensor biểu diễn trung bình mất mát -1T Σ_t log pi_t · (G_t - b_t).

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ả.