Bỏ qua

Học Tăng Cường (Reinforcement Learning)

Học tăng cường huấn luyện các tác tử đưa ra quyết định tuần tự bằng cách cực đại hóa phần thưởng tích lũy thông qua thử và sai. File này bao gồm MDP, hàm giá trị, phương trình Bellman, Q-learning, policy gradient, phương pháp actor-critic, PPO và RLHF — khung làm việc đứng sau các tác tử chơi game và sự căn chỉnh mô hình ngôn ngữ.

  • Học có giám sát cần dữ liệu có nhãn. Học không giám sát tìm các mẫu trong dữ liệu không nhãn. Học tăng cường (Reinforcement learning - RL) khác cả hai: một tác tử (agent) học bằng cách tương tác với môi trường, thực hiện hành động, và nhận phần thưởng. Không có các nhãn đúng; tác tử phải tự khám phá ra hành vi tốt thông qua thử và sai.

  • Hãy nghĩ về việc dạy một con chó một trò mới. Bạn không đưa ra cho nó một bộ dữ liệu về các hành vi đúng. Thay vào đó, nó thử các việc, bạn cho thưởng khi nó làm đúng, và theo thời gian nó tìm ra bạn muốn gì. RL chính thức hóa quá trình này.

  • Thiết lập RL có năm thành phần cốt lõi. Tác tử (Agent) là người học và ra quyết định. Môi trường (Environment) là mọi thứ bên ngoài tác tử mà nó tương tác. Tại mỗi bước thời gian, tác tử quan sát một trạng thái (state) \(s_t\), chọn một hành động (action) \(a_t\), nhận một phần thưởng (reward) \(r_t\), và chuyển sang trạng thái mới \(s_{t+1}\). Mục tiêu của tác tử là cực đại hóa tổng phần thưởng thu được theo thời gian.

Vòng lặp tác tử-môi trường: tác tử quan sát trạng thái, thực hiện hành động, nhận phần thưởng, môi trường chuyển sang trạng thái mới

  • Chính sách (Policy) \(\pi\) là chiến lược của tác tử: một ánh xạ từ trạng thái sang hành động. Một chính sách tất định (deterministic) cho một hành động mỗi trạng thái: \(a = \pi(s)\). Một chính sách ngẫu nhiên (stochastic) cho một phân bố xác suất trên các hành động: \(\pi(a \mid s)\). Mục tiêu của RL là tìm chính sách tối ưu, chính sách cực đại hóa phần thưởng tích lũy kỳ vọng.

  • Khung toán học cho RL là Quá trình Quyết định Markov (Markov Decision Process - MDP), được định nghĩa bởi một bộ \((S, A, P, R, \gamma)\): một tập các trạng thái \(S\), một tập các hành động \(A\), các xác suất chuyển tiếp \(P(s' \mid s, a)\), một hàm phần thưởng \(R(s, a)\), và một hệ số chiết khấu \(\gamma\).

  • Tính chất Markov (từ chương 05) nói rằng tương lai chỉ phụ thuộc vào trạng thái hiện tại, không phải lịch sử bạn đã đến đó như thế nào: \(P(s_{t+1} \mid s_t, a_t, s_{t-1}, \ldots) = P(s_{t+1} \mid s_t, a_t)\). Điều này có nghĩa trạng thái chứa mọi thông tin cần để đưa ra quyết định.

  • Hệ số chiết khấu (Discount factor) \(\gamma \in [0, 1)\) quyết định tác tử quan tâm đến các phần thưởng tương lai so với tức thời như thế nào. Lợi tức chiết khấu (discounted return) từ thời điểm \(t\) là:

\[G_t = r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \cdots = \sum_{k=0}^{\infty} \gamma^k r_{t+k}\]
  • Với \(\gamma = 0\), tác tử hoàn toàn thiển cận, chỉ quan tâm đến phần thưởng kế tiếp. Với \(\gamma\) gần 1, tác tử có tầm nhìn xa. Hệ số chiết khấu cũng đảm bảo tổng hội tụ (nếu các phần thưởng bị chặn), điều quan trọng cho tính xác định toán học.

  • Các hàm giá trị (Value functions) ước lượng tốt như thế nào khi ở một trạng thái (hoặc thực hiện một hành động trong một trạng thái). Hàm giá trị trạng thái (State-value function) \(V^\pi(s)\) là lợi tức kỳ vọng bắt đầu từ trạng thái \(s\) và đi theo chính sách \(\pi\):

\[V^\pi(s) = \mathbb{E}_\pi \left[ G_t \mid s_t = s \right]\]
  • Hàm giá trị hành động (Action-value function) \(Q^\pi(s, a)\) là lợi tức kỳ vọng bắt đầu từ trạng thái \(s\), thực hiện hành động \(a\), và sau đó đi theo \(\pi\):
\[Q^\pi(s, a) = \mathbb{E}_\pi \left[ G_t \mid s_t = s, a_t = a \right]\]
  • Mối quan hệ: \(V^\pi(s) = \sum_a \pi(a \mid s) \, Q^\pi(s, a)\). Giá trị trạng thái là trung bình của các giá trị hành động, có trọng số theo chính sách.

  • Phương trình Bellman biểu diễn một quan hệ hồi quy: giá trị của một trạng thái bằng phần thưởng tức thời cộng với giá trị chiết khấu của trạng thái tiếp theo. Với hàm giá trị trạng thái:

\[V^\pi(s) = \sum_a \pi(a \mid s) \sum_{s'} P(s' \mid s, a) \left[ R(s, a) + \gamma \, V^\pi(s') \right]\]
  • Với hàm giá trị tối ưu \(V^{*}(s)\), tác tử luôn chọn hành động tốt nhất:
\[V^{*}(s) = \max_a \sum_{s'} P(s' \mid s, a) \left[ R(s, a) + \gamma \, V^{*}(s') \right]\]
  • Tương tự, phương trình Bellman tối ưu cho \(Q^{*}\):
\[Q^{*}(s, a) = \sum_{s'} P(s' \mid s, a) \left[ R(s, a) + \gamma \max_{a'} Q^{*}(s', a') \right]\]
  • Một khi bạn có \(Q^{*}\), chính sách tối ưu rất đơn giản: luôn chọn hành động có Q-value cao nhất: \(\pi^{*}(s) = \arg\max_a Q^{*}(s, a)\).

  • Các phương pháp Quy hoạch động (Dynamic programming) giải MDP khi bạn biết các xác suất chuyển tiếp và phần thưởng (mô hình đầy đủ). Đánh giá chính sách (Policy evaluation) tính \(V^\pi\) cho một chính sách cho trước bằng cách áp dụng lặp đi lặp lại phương trình Bellman cho đến khi hội tụ. Cải tiến chính sách (Policy improvement) lấy hàm giá trị và xây dựng một chính sách tốt hơn bằng cách hành động tham lam: \(\pi'(s) = \arg\max_a \sum_{s'} P(s' \mid s, a)[R(s,a) + \gamma V^\pi(s')]\).

  • Lặp chính sách (Policy iteration) luân phiên giữa đánh giá và cải tiến cho đến khi chính sách dừng thay đổi. Nó được đảm bảo hội tụ đến chính sách tối ưu.

  • Lặp giá trị (Value iteration) kết hợp cả hai bước thành một: nó lặp đi lặp lại áp dụng phương trình Bellman tối ưu cho đến khi \(V^{*}\) hội tụ, sau đó trích xuất chính sách.

\[V(s) \leftarrow \max_a \sum_{s'} P(s' \mid s, a) \left[ R(s, a) + \gamma \, V(s') \right]\]
  • Quy hoạch động yêu cầu biết \(P(s' \mid s, a)\), điều thường phi thực tế. Trong hầu hết các bài toán thực tế, tác tử không biết động lực học (dynamics) của môi trường; nó chỉ có thể tương tác với môi trường. Đây là lúc các phương pháp không mô hình (model-free) ra đời.

  • Học chênh lệch thời gian (Temporal Difference - TD learning) học từ kinh nghiệm mà không biết mô hình. Ý tưởng chính là bootstrapping: thay vì đợi đến cuối tập để tính lợi tức thực tế \(G_t\), bạn ước lượng nó bằng hàm giá trị hiện tại:

\[V(s_t) \leftarrow V(s_t) + \alpha \left[ r_t + \gamma \, V(s_{t+1}) - V(s_t) \right]\]
  • Phần trong ngoặc là sai số TD (TD error): hiệu số giữa mục tiêu TD (\(r_t + \gamma V(s_{t+1})\)) và ước lượng hiện tại \(V(s_t)\). Nếu sai số TD dương, trạng thái tốt hơn kỳ vọng, do đó chúng ta tăng giá trị của nó. Nếu âm, chúng ta giảm nó.

Chuyển tiếp trạng thái thể hiện mục tiêu TD: giá trị hiện tại, phần thưởng, và giá trị kế tiếp bootstrap với công thức cập nhật

  • TD learning cập nhật sau mỗi bước đơn lẻ (không sau cả tập), làm cho nó hiệu quả hơn nhiều so với các phương pháp Monte Carlo. Nó cũng hoạt động trong các môi trường tiếp diễn (không theo tập).

  • SARSA (State-Action-Reward-State-Action) là TD learning áp dụng cho Q-values. Tác tử thực hiện hành động \(a\) trong trạng thái \(s\), quan sát phần thưởng \(r\) và trạng thái tiếp theo \(s'\), sau đó chọn hành động tiếp theo \(a'\) theo chính sách của nó:

\[Q(s, a) \leftarrow Q(s, a) + \alpha \left[ r + \gamma \, Q(s', a') - Q(s, a) \right]\]
  • SARSA là on-policy: nó cập nhật bằng cách dùng hành động mà tác tử thực sự thực hiện, bao gồm cả thăm dò. Điều này làm SARSA bảo thủ hơn; nó học một chính sách tính đến nhiễu thăm dò của chính nó.

  • Q-learning là thuật toán RL nổi tiếng nhất. Nó giống SARSA, nhưng thay vì dùng hành động tác tử thực sự thực hiện, nó dùng hành động tốt nhất có thể:

\[Q(s, a) \leftarrow Q(s, a) + \alpha \left[ r + \gamma \max_{a'} Q(s', a') - Q(s, a) \right]\]
  • Q-learning là off-policy: nó học các Q-value tối ưu bất kể chính sách đang được theo dõi. Tác tử có thể thăm dò ngẫu nhiên trong khi vẫn học được các giá trị hành động tối ưu. Điều này làm Q-learning mạnh mẽ hơn và thường hội tụ nhanh hơn, nhưng nó có thể ước lượng quá cao (overestimate) các giá trị.

  • Khám phá vs khai thác (Exploration vs exploitation) là vấn đề nan giải cơ bản: tác tử nên khai thác những gì nó đã biết (chọn hành động có giá trị ước lượng cao nhất) hay khám phá các hành động chưa biết (mà có thể hóa ra là tốt hơn)?

  • Chiến lược đơn giản nhất là epsilon-tham lam (epsilon-greedy): với xác suất \(\epsilon\), thực hiện một hành động ngẫu nhiên (khám phá); với xác suất \(1 - \epsilon\), thực hiện hành động tham lam (khai thác). Một lịch trình phổ biến bắt đầu với \(\epsilon\) cao (nhiều khám phá) và suy giảm nó theo thời gian.

  • Các phương pháp dạng bảng (lưu một giá trị cho mỗi cặp trạng thái-hành động trong một bảng) hoạt động cho các không gian trạng thái nhỏ, rời rạc. Với các không gian trạng thái lớn hoặc liên tục, bạn cần xấp xỉ hàm. Mạng Q sâu (Deep Q-Networks - DQN) sử dụng một mạng nơ-ron để xấp xỉ \(Q(s, a; \theta)\), trong đó \(\theta\) là các trọng số mạng.

  • DQN giới thiệu hai kỹ thuật ổn định quan trọng. Phát lại kinh nghiệm (Experience replay): thay vì học từ các chuyển tiếp liên tiếp (có tương quan cao), lưu các chuyển tiếp vào một bộ đệm phát lại (replay buffer) và lấy mẫu các mini-batch ngẫu nhiên để huấn luyện. Cách này phá vỡ tương quan và tái sử dụng dữ liệu hiệu quả.

  • Mạng mục tiêu (Target network): sử dụng một bản sao riêng, được cập nhật chậm của mạng để tính các mục tiêu TD. Nếu không có điều này, mục tiêu di chuyển mỗi lần bạn cập nhật mạng, tạo ra sự bất ổn định kiểu "đuổi theo đuôi của chính mình". Mạng mục tiêu được cập nhật định kỳ (hard update mỗi \(N\) bước) hoặc liên tục (soft update: \(\theta^{-} \leftarrow \tau\theta + (1-\tau)\theta^{-}\)).

  • Mất mát DQN chỉ là MSE giữa Q-value dự đoán và mục tiêu TD:

\[\mathcal{L}(\theta) = \mathbb{E} \left[ \left( r + \gamma \max_{a'} Q(s', a'; \theta^{-}) - Q(s, a; \theta) \right)^2 \right]\]
  • Tất cả các phương pháp cho đến giờ đều học các hàm giá trị và suy ra chính sách từ chúng. Các phương pháp Policy gradient có cách tiếp cận khác: chúng tham số hóa trực tiếp chính sách \(\pi(a \mid s; \theta)\) và tối ưu hóa nó bằng hạ gradient (lên dốc) trên lợi tức kỳ vọng.

  • Định lý policy gradient cho gradient của lợi tức kỳ vọng đối với các tham số chính sách:

\[\nabla_\theta J(\theta) = \mathbb{E}_\pi \left[ \nabla_\theta \log \pi(a \mid s; \theta) \cdot G_t \right]\]
  • Nó nói: tăng xác suất của các hành động dẫn đến lợi tức cao, giảm xác suất của các hành động dẫn đến lợi tức thấp. Gradient log-xác suất cho hướng thay đổi chính sách, và \(G_t\) tỷ lệ hóa mức thay đổi.

  • REINFORCE là thuật toán policy gradient đơn giản nhất. Chạy một tập, tính các lợi tức \(G_t\) cho mỗi bước, và cập nhật:

\[\theta \leftarrow \theta + \alpha \, \nabla_\theta \log \pi(a_t \mid s_t; \theta) \cdot G_t\]
  • REINFORCE có phương sai cao vì \(G_t\) là ước lượng nhiễu từ một mẫu đơn của lợi tức kỳ vọng. Một sửa chữa phổ biến là trừ đi một đường cơ sở (baseline) (thường là lợi tức trung bình hoặc một hàm giá trị đã học) để giảm phương sai mà không đưa vào độ chệch:
\[\theta \leftarrow \theta + \alpha \, \nabla_\theta \log \pi(a_t \mid s_t; \theta) \cdot (G_t - b)\]
  • Các phương pháp Actor-Critic sử dụng hai mạng. Actor là chính sách \(\pi(a \mid s; \theta)\). Critic là một hàm giá trị \(V(s; \phi)\) đóng vai trò là đường cơ sở. Lợi thế \(A_t = r_t + \gamma V(s_{t+1}) - V(s_t)\) thay thế \(G_t - b\):
\[\theta \leftarrow \theta + \alpha \, \nabla_\theta \log \pi(a_t \mid s_t; \theta) \cdot A_t\]
  • Critic được cập nhật bằng cách tối thiểu hóa sai số TD, giống các phương pháp dựa trên giá trị. Actor được cập nhật bằng policy gradient, với ước lượng lợi thế của critic giúp giảm phương sai. Đây là tinh hoa của cả hai thế giới.

Kiến trúc hai đầu: actor xuất ra các xác suất hành động, critic xuất ra ước lượng giá trị, tín hiệu lợi thế hướng dẫn cập nhật actor

  • PPO (Proximal Policy Optimization) là thuật toán policy gradient được sử dụng rộng rãi nhất trong thực tế. Nó giải quyết một vấn đề then chốt: nếu một cập nhật chính sách quá lớn, hiệu suất có thể sụp đổ thảm hại.

  • PPO sử dụng một mục tiêu cắt (clipped surrogate objective). Gọi \(r_t(\theta) = \frac{\pi(a_t | s_t; \theta)}{\pi(a_t | s_t; \theta_{\text{old}})}\) là tỷ lệ xác suất giữa chính sách mới và cũ. Mất mát là:

\[\mathcal{L}^{\text{CLIP}}(\theta) = \mathbb{E} \left[ \min\!\left( r_t(\theta) A_t, \; \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t \right) \right]\]
  • Việc cắt (thường \(\epsilon = 0.2\)) ngăn tỷ lệ di chuyển quá xa khỏi 1, giữ cho các cập nhật nhỏ và ổn định. Nếu lợi thế dương (hành động tốt), tỷ lệ được chặn trên tại \(1 + \epsilon\). Nếu âm (hành động xấu), tỷ lệ được chặn dưới tại \(1 - \epsilon\). Cách này đơn giản và ổn định hơn các phương pháp trust-region trước đó (TRPO).

  • PPO là thuật toán đã được dùng để huấn luyện các mô hình kiểu ChatGPT thông qua RLHF (Reinforcement Learning from Human Feedback - Học Tăng Cường từ Phản Hồi Con Người). Trong RLHF, một mô hình phần thưởng (reward model) được huấn luyện trên dữ liệu ưu tiên của con người (trong hai đầu ra, con người thích cái nào hơn?), và sau đó PPO tối ưu hóa chính sách của mô hình ngôn ngữ để cực đại hóa phần thưởng đã học này.

  • DPO (Direct Preference Optimization - Tối ưu hóa Ưu tiên Trực tiếp) đơn giản hóa RLHF bằng cách loại bỏ hoàn toàn mô hình phần thưởng. Thay vì huấn luyện một mô hình phần thưởng và sau đó chạy RL, DPO suy ra một mất mát dạng đóng tối ưu hóa trực tiếp chính sách từ dữ liệu ưu tiên:

\[\mathcal{L}_{\text{DPO}}(\theta) = -\mathbb{E} \left[ \log \sigma\!\left( \beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)} \right) \right]\]
  • Ở đây \(y_w\) là phản hồi được ưa thích (thắng) và \(y_l\) là phản hồi không được ưa thích (thua). DPO tăng xác suất tương đối của các đầu ra được ưa thích và đơn giản hơn nhiều để cài đặt so với RLHF dựa trên PPO.

  • Hai phân biệt quan trọng trong các thuật toán RL. On-policy vs off-policy: phương pháp on-policy (SARSA, PPO) học từ dữ liệu do chính sách hiện tại tạo ra; phương pháp off-policy (Q-learning, DQN) có thể học từ dữ liệu do bất kỳ chính sách nào tạo ra. Phương pháp off-policy tiết kiệm mẫu hơn (chúng tái sử dụng dữ liệu cũ) nhưng có thể kém ổn định hơn.

  • Dựa trên mô hình (Model-based) vs không mô hình (model-free): phương pháp model-free (mọi thứ đã thảo luận cho đến nay) học các giá trị hoặc chính sách trực tiếp từ kinh nghiệm. Phương pháp model-based học một mô hình của môi trường (\(P(s' \mid s, a)\)\(R(s, a)\)) và dùng nó cho lập kế hoạch (tưởng tượng các quỹ đạo tương lai mà không thực sự thực hiện hành động). Phương pháp model-based tiết kiệm mẫu hơn nhưng thêm độ phức tạp của việc học một mô hình chính xác.

  • Tóm tắt bối cảnh RL:

Phương pháp Loại Ý tưởng chính Điểm mạnh
Lặp giá trị DP, model-based Tối ưu Bellman Nghiệm chính xác (MDP nhỏ)
SARSA TD, on-policy Học Q on-policy Bảo thủ, an toàn
Q-Learning TD, off-policy Học Q*, mục tiêu tham lam Đơn giản, hiệu quả
DQN Sâu, off-policy Neural Q + replay + target net Co giãn đến trạng thái chiều cao
REINFORCE Policy gradient Gradient log-prob * lợi tức Tối ưu chính sách đơn giản
Actor-Critic PG + giá trị Actor + critic cho phương sai thấp Thiết thực và linh hoạt
PPO PG, cắt Ổn định kiểu trust-region Tiêu chuẩn công nghiệp
DPO Ưu tiên trực tiếp Bỏ qua mô hình phần thưởng RLHF đơn giản hơn

Bài tập lập trình (dùng CoLab hoặc notebook)

  1. Cài đặt lặp giá trị (value iteration) cho một gridworld đơn giản. Tính hàm giá trị tối ưu và trích xuất chính sách tối ưu. Trực quan hóa cả hai dưới dạng bản đồ nhiệt và đồ thị mũi tên.

    import jax.numpy as jnp
    import matplotlib.pyplot as plt
    
    # 4x4 gridworld: goal at (3,3), reward -1 per step, 0 at goal
    grid_size = 4
    gamma = 0.99
    goal = (3, 3)
    
    # Actions: up, down, left, right
    actions = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    action_names = ['up', 'down', 'left', 'right']
    action_arrows = ['\u2191', '\u2193', '\u2190', '\u2192']
    
    def step(s, a):
        """Deterministic transition."""
        ns = (max(0, min(grid_size-1, s[0]+a[0])),
              max(0, min(grid_size-1, s[1]+a[1])))
        return ns
    
    # Value iteration
    V = jnp.zeros((grid_size, grid_size))
    for iteration in range(100):
        V_new = jnp.array(V)
        for i in range(grid_size):
            for j in range(grid_size):
                if (i, j) == goal:
                    continue
                values = []
                for a in actions:
                    ns = step((i, j), a)
                    values.append(-1 + gamma * float(V[ns[0], ns[1]]))
                V_new = V_new.at[i, j].set(max(values))
        if jnp.max(jnp.abs(V_new - V)) < 1e-6:
            print(f"Converged in {iteration+1} iterations")
            break
        V = V_new
    
    # Extract policy
    policy = [['' for _ in range(grid_size)] for _ in range(grid_size)]
    for i in range(grid_size):
        for j in range(grid_size):
            if (i, j) == goal:
                policy[i][j] = 'G'
                continue
            best_a = max(range(4), key=lambda a: -1 + gamma * float(V[step((i,j), actions[a])[0], step((i,j), actions[a])[1]]))
            policy[i][j] = action_arrows[best_a]
    
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    im = axes[0].imshow(V, cmap='YlOrRd_r')
    axes[0].set_title("Optimal Value Function")
    for i in range(grid_size):
        for j in range(grid_size):
            axes[0].text(j, i, f"{V[i,j]:.1f}", ha='center', va='center', fontsize=10)
    plt.colorbar(im, ax=axes[0])
    
    axes[1].imshow(jnp.ones((grid_size, grid_size)), cmap='Greys', vmin=0, vmax=2)
    axes[1].set_title("Optimal Policy")
    for i in range(grid_size):
        for j in range(grid_size):
            axes[1].text(j, i, policy[i][j], ha='center', va='center', fontsize=18)
    plt.tight_layout(); plt.show()
    

  2. Cài đặt Q-learning dạng bảng trên một gridworld đơn giản. Huấn luyện tác tử, vẽ đường cong học và hiển thị các Q-values đã học.

    import jax
    import jax.numpy as jnp
    import matplotlib.pyplot as plt
    
    grid_size = 5
    goal = (4, 4)
    actions = [(-1,0), (1,0), (0,-1), (0,1)]
    
    # Q-table
    Q = {}
    for i in range(grid_size):
        for j in range(grid_size):
            Q[(i,j)] = [0.0] * 4
    
    alpha = 0.1
    gamma = 0.95
    epsilon = 1.0
    epsilon_decay = 0.995
    min_epsilon = 0.01
    
    def step(s, a_idx):
        a = actions[a_idx]
        ns = (max(0, min(grid_size-1, s[0]+a[0])),
              max(0, min(grid_size-1, s[1]+a[1])))
        r = 0.0 if ns == goal else -1.0
        done = ns == goal
        return ns, r, done
    
    key = jax.random.PRNGKey(42)
    rewards_per_episode = []
    
    for ep in range(500):
        s = (0, 0)
        total_reward = 0
        for _ in range(100):
            key, subkey = jax.random.split(key)
            if float(jax.random.uniform(subkey)) < epsilon:
                key, subkey = jax.random.split(key)
                a = int(jax.random.randint(subkey, (), 0, 4))
            else:
                a = max(range(4), key=lambda i: Q[s][i])
    
            ns, r, done = step(s, a)
            total_reward += r
            # Q-learning update
            Q[s][a] += alpha * (r + gamma * max(Q[ns]) - Q[s][a])
            s = ns
            if done:
                break
        rewards_per_episode.append(total_reward)
        epsilon = max(min_epsilon, epsilon * epsilon_decay)
    
    plt.figure(figsize=(8, 4))
    # Smooth the curve
    window = 20
    smoothed = [sum(rewards_per_episode[max(0,i-window):i+1])/min(i+1, window)
                for i in range(len(rewards_per_episode))]
    plt.plot(smoothed, color='#3498db', linewidth=1.5)
    plt.xlabel("Episode"); plt.ylabel("Total Reward (smoothed)")
    plt.title("Q-Learning on Gridworld")
    plt.grid(alpha=0.3); plt.show()
    
    # Show learned policy
    arrow = ['\u2191', '\u2193', '\u2190', '\u2192']
    print("Learned policy:")
    for i in range(grid_size):
        row = ""
        for j in range(grid_size):
            if (i,j) == goal:
                row += " G "
            else:
                row += f" {arrow[max(range(4), key=lambda a: Q[(i,j)][a])]} "
        print(row)
    

  3. Cài đặt REINFORCE trên một bài toán multi-armed bandit. Chỉ ra cách chính sách tiến hóa theo thời gian để ưu tiên nhánh tốt nhất.

    import jax
    import jax.numpy as jnp
    import matplotlib.pyplot as plt
    
    # 5-armed bandit with different expected rewards
    true_rewards = jnp.array([0.2, 0.5, 0.8, 0.3, 0.1])
    n_arms = len(true_rewards)
    
    # Policy: softmax over logits
    logits = jnp.zeros(n_arms)
    lr = 0.1
    key = jax.random.PRNGKey(42)
    
    policy_history = []
    reward_history = []
    
    for step in range(2000):
        probs = jax.nn.softmax(logits)
        policy_history.append(probs)
    
        # Sample action
        key, subkey = jax.random.split(key)
        action = jax.random.choice(subkey, n_arms, p=probs)
    
        # Get reward (Bernoulli)
        key, subkey = jax.random.split(key)
        reward = float(jax.random.uniform(subkey) < true_rewards[action])
        reward_history.append(reward)
    
        # REINFORCE update
        # grad log pi(a) = e_a - probs (for softmax parameterisation)
        grad_log_pi = -probs.at[action].add(1.0)  # one-hot(a) - probs
        logits = logits + lr * reward * grad_log_pi
    
    policy_history = jnp.stack(policy_history)
    
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    colors = ['#3498db', '#e74c3c', '#27ae60', '#9b59b6', '#f39c12']
    for i in range(n_arms):
        axes[0].plot(policy_history[:, i], color=colors[i],
                     label=f'Arm {i} (true={true_rewards[i]:.1f})', linewidth=1.5)
    axes[0].set_xlabel("Step"); axes[0].set_ylabel("P(arm)")
    axes[0].set_title("Policy Evolution (REINFORCE)")
    axes[0].legend(fontsize=8); axes[0].grid(alpha=0.3)
    
    # Smoothed reward
    window = 50
    smoothed = [sum(reward_history[max(0,i-window):i+1])/min(i+1,window)
                for i in range(len(reward_history))]
    axes[1].plot(smoothed, color='#27ae60', linewidth=1.5)
    axes[1].axhline(y=0.8, color='#e74c3c', linestyle='--', alpha=0.5, label='Best arm')
    axes[1].set_xlabel("Step"); axes[1].set_ylabel("Avg Reward")
    axes[1].set_title("Reward Over Time"); axes[1].legend()
    axes[1].grid(alpha=0.3)
    plt.tight_layout(); plt.show()