Log-derivative trick: vì sao đạo hàm của log xác suất lại hữu ích

$$ \nabla_\theta \log p_\theta(x) = \frac{1}{p_\theta(x)},\nabla_\theta p_\theta(x) \quad\Longrightarrow\quad \nabla_\theta p_\theta(x) = p_\theta(x),\nabla_\theta \log p_\theta(x) $$

Hai lưu ý nhỏ nhưng hay bị bỏ. Một, log ở đây hiểu là log tự nhiên ($\ln$); nếu dùng cơ số khác thì hai vế lệch nhau một hằng số, nên trong tài liệu RL luôn ngầm hiểu log là $\ln$. Hai, $p_\theta(x)$ là xác suất của một điểm dữ liệu hoặc một hành động cụ thể, không phải xác suất của cả phân phối. Trên mỗi mẫu, $\nabla \log p$ được gọi là score function (hàm điểm).

Bài tập 1 và gợi ý

Cho policy Bernoulli $p = \sigma(\theta)$, tức $p_\theta(x{=}1) = \sigma(\theta)$ và $p_\theta(x{=}0) = 1 - \sigma(\theta)$. Tính $\nabla_\theta \log p_\theta(x{=}1)$ và $\nabla_\theta \log p_\theta(x{=}0)$.

Gợi ý: dùng $\frac{d}{d\theta}\sigma(\theta) = \sigma(\theta)(1-\sigma(\theta))$, rồi áp công thức vừa chứng minh.

Đáp án: $\nabla_\theta \log p_\theta(x{=}1) = 1 - \sigma(\theta)$ (cũng bằng $\sigma(-\theta)$); $\nabla_\theta \log p_\theta(x{=}0) = -\sigma(\theta)$.

Đưa đạo hàm vào trong dấu kỳ vọng

Thay đẳng thức vừa chứng minh vào biểu thức của $\nabla J$:

$$ \nabla_\theta J(\theta) = \sum_x p_\theta(x),\nabla_\theta \log p_\theta(x), f(x) = \mathbb{E}{x \sim p\theta}\left[\nabla_\theta \log p_\theta(x), f(x)\right] $$

Đây chính là toàn bộ giá trị của thủ thuật: đạo hàm của một kỳ vọng đã trở thành kỳ vọng của một đại lượng, và đại lượng đó ta tính được cho từng mẫu. Vì vậy ước lượng Monte Carlo chỉ còn là một phép trung bình:

$$ \nabla_\theta J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N}\nabla_\theta \log p_\theta(x_i), f(x_i), \qquad x_i \sim p_\theta $$

Hai hệ quả rất thực dụng cho đồ án. Một, $f$ không cần khả vi, bạn chỉ cần tính được giá trị phần thưởng; môi trường hoàn toàn có thể là hộp đen. Hai, ta không cần đạo hàm xuyên qua môi trường, chỉ cần đạo hàm xuyên qua policy của mình.

Một tính chất đẹp đi kèm: vì $\sum_x p_\theta(x) = 1$ với mọi $\theta$ nên $\nabla_\theta \sum_x p_\theta(x) = 0$, kéo theo $\mathbb{E}{x \sim p\theta}[\nabla_\theta \log p_\theta(x)] = 0$. Hệ quả là cộng thêm một hằng số $c$ vào $f$ không làm đổi kỳ vọng của gradient. Bạn sẽ khai thác điểm này khi làm baseline ở bài sau.

Bài tập 2 và gợi ý

Chứng minh $\mathbb{E}{x \sim p\theta}[\nabla_\theta \log p_\theta(x)] = 0$ bằng cách dùng lại chính thủ thuật này.

Gợi ý: viết kỳ vọng thành tổng, thay $\nabla \log p$ bằng $\frac{1}{p}\nabla p$, rồi dùng $\nabla_\theta \sum_x p_\theta(x) = \nabla_\theta 1 = 0$.

Đáp án: $\sum_x p_\theta(x)\frac{\nabla p_\theta(x)}{p_\theta(x)} = \sum_x \nabla_\theta p_\theta(x) = \nabla_\theta \sum_x p_\theta(x) = 0$.

Dạng bạn sẽ gõ trong đồ án

Trong REINFORCE, $x$ chính là hành động $a$ và $p_\theta$ là policy $\pi_\theta(a \mid s)$. Với policy softmax trên logits $z$, score có dạng rất gọn:

$$ \nabla_z \log \pi_\theta(a \mid s) = \mathbf{1}a - \pi\theta(\cdot \mid s) $$

tức vector one-hot của hành động vừa chọn trừ đi vector xác suất. Khung xương trong code sẽ giống thế này:

# bản nháp ý tưởng, bài 2.1 sẽ dựng đầy đủ
log_probs, rewards = [], []
for t in range(T):
    a, logp = policy.act(s)      # lấy mẫu a ~ π(·|s), trả về log π
    s, r, done = env.step(a)
    log_probs.append(logp)
    rewards.append(r)
# gradient của quỹ đạo: sum_t ∇log π(a_t|s_t) * G_t

Bài tập 3 và đáp án

Policy có hai hành động với $\pi = (0.8, 0.2)$. Tính $\nabla_z \log \pi(a)$ khi $a$ là hành động thứ nhất và khi $a$ là hành động thứ hai.

Đáp án: với $a = 1$, vector là $(1-0.8,; 0-0.2) = (0.2,,-0.2)$. Với $a = 2$, vector là $(-0.8,,0.8)$. Nhận xét: hành động càng ít khả năng thì gradient càng lớn, nên bước cập nhật càng mạnh khi nó bất ngờ xảy ra.

Bốn cái bẫy khi dùng log-derivative trick

• Lấy log sai đối tượng. Log áp lên xác suất $p_\theta(x)$, không phải lên phần thưởng $f(x)$ và cũng không phải lên $J$. Nhầm chỗ này là công thức sai ngay từ dòng đầu.
• Quên điều kiện xác suất dương. $\log 0$ không xác định, trong code là -inf và lan ra toàn bộ gradient. Hãy dùng log_softmax thay vì log(softmax(...)) để ổn định số học (numerical stability).
• Tưởng thủ thuật này khử variance. Nó chỉ đổi chỗ đạo hàm, không làm gradient bớt nhiễu. Variance của ước lượng này thực tế rất lớn, đó là lý do bài sau cần baseline.
• **Nhân sai chiều. Nhớ rằng $f(x)$ là một số vô hướng nhân vào cả vector gradient. Nếu bạn lấy trung bình cộng các $\nabla_\theta \log p_\theta(x_i)$ mà bỏ $f(x_i)$ ra ngoài, bạn đang đẩy xác suất của mọi hành động lên như nhau, kể cả hành động tệ.

Quy tắc ngón tay cái: mỗi khi bạn thấy mình loay hoay đạo hàm một hàm phần thưởng không khả vi theo $\theta$, hãy dừng lại và tự hỏi: "Mình có viết được đại lượng này thành kỳ vọng theo một phân phối khả vi theo $\theta$ không?" Nếu có, log-derivative trick lo phần còn lại.

Tự kiểm tra nhanh bằng điền khuyết

Đẳng thức lõi của log-derivative trick là ∇_θ p_θ(x) = _____. Đẳng thức đó đúng vì theo quy tắc chuỗi cho hàm log, ta có ∇_θ log p_θ(x) = _____, chỉ cần _____. Đại lượng ∇_θ log p_θ(x) thường được gọi là _____. Nhờ trick, gradient của mục tiêu trở thành một _____ theo chính phân phối p_θ, nên có thể ước lượng bằng _____.

Bộ thẻ ôn tập log-derivative trick

Front: Phát biểu log-derivative trick Back: ∇_θ p_θ(x) = p_θ(x) ∇_θ log p_θ(x).

Front: Vì sao đẳng thức ∇p = p ∇log p đúng? Back: Vì ∇log p = ∇p / p, nhân hai vế với p là xong; đây chỉ là quy tắc chuỗi cho hàm log.

Front: Điều kiện để lấy log hợp lệ là gì? Back: p_θ(x) > 0 với mọi x có thể lấy mẫu; log 0 không xác định.

Front: ∇_θ log p_θ(x) có tên gọi là gì? Back: Hàm score (score function) của phân phối.

Front: Nhờ trick, gradient của J(θ) = E_{xp_θ}[f(x)] viết lại thế nào? Back: ∇J(θ) = E_{xp_θ}[ f(x) ∇_θ log p_θ(x) ].

Front: Vì sao dạng mới lấy mẫu được còn dạng cũ thì không? Back: Dạng mới không còn đạo hàm của f theo θ; chỉ cần mẫu x ~ p_θ và tính ∇log p_θ(x).

Front: Ước lượng Monte Carlo của ∇J trông ra sao? Back: Trung bình cộng của f(x_i) · ∇_θ log p_θ(x_i) trên N mẫu.

Front: Với policy softmax, ∇_θ log π_θ(a|s) bằng gì? Back: Vector one-hot của hành động a trừ đi vector xác suất π_θ(·|s).

Front: Log-derivative trick có làm giảm variance không? Back: Không. Nó chỉ đổi chỗ đạo hàm; variance của ước lượng vẫn rất cao, cần baseline ở bài sau.

Front: Trong REINFORCE, f(x) tương ứng với đại lượng nào? Back: Phần thưởng của quỹ đạo, thường là return G_t gắn với cặp (s_t, a_t).

Front: Tên gọi khác của log-derivative trick? Back: Likelihood ratio trick, score function estimator, hay REINFORCE trick.

Front: Vì sao không đạo hàm trực tiếp phần thưởng f(x) theo θ? Back: Vì f thường không phụ thuộc tường minh vào θ và không khả vi; trick chuyển toàn bộ đạo hàm sang p_θ.

Mang gì sang bài sau

Ba điều cần nhớ: (1) $\nabla_\theta p = p ,\nabla_\theta \log p$ chỉ là quy tắc chuỗi viết lại, (2) nó biến đạo hàm của một kỳ vọng thành kỳ vọng của một tích, và (3) nhờ vậy ta ước lượng được gradient chỉ bằng cách lấy mẫu và cộng dồn, đây chính là hạt nhân tính được của Policy Gradient. Sang bài tiếp theo, ta sẽ thay $f(x)$ bằng return của môi trường và ghép mọi thứ thành công thức cập nhật policy hoàn chỉnh, bước đầu tiên để viết vòng lặp REINFORCE trong đồ án.