Tự xây dựng mô hình thế giới (World Model) đầu tiên của bạn bằng Python

Mô hình thế giới (world model) là một trong những ý tưởng thú vị nhất trong học tăng cường hiện đại: thay vì chỉ phản ứng với môi trường, tác nhân (agent) học cách mô phỏng thế giới trong đầu, rồi "mơ mộng" về tương lai để lập kế hoạch. Bài viết này sẽ hướng dẫn bạn xây dựng một world model đơn giản từ đầu, cho nó tự diễn giải qua trò chơi CartPole, và quan trọng nhất — đo lường chính xác thời điểm ảo giác sụp đổ.

World Model là gì và tại sao lại quan trọng?

Hãy tưởng tượng bạn chơi một trò chơi điện tử. Bạn không cần nhìn vào màn hình để biết rằng nếu bấm nút nhảy, nhân vật sẽ bay lên. Bạn có một mô hình tinh thần về luật chơi trong đầu. Đó chính là world model.

Trong học máy, world model là một mạng nơ-ron học cách dự đoán:

• Trạng thái tiếp theo của môi trường khi biết trạng thái hiện tại và hành động
• Phần thưởng mà tác nhân sẽ nhận được
• Trạng thái kết thúc (liệu ván chơi có kết thúc không)

Khi đã có world model, tác nhân có thể "suy nghĩ" mà không cần tương tác thật với môi trường — tiết kiệm thời gian, chi phí và cả những va chạm nguy hiểm trong thế giới thực.

Các hệ thống như Dreamer, World Models của Ha & Schmidhuber, và gần đây là các mô hình nền tảng cho robot đều dựa trên ý tưởng cốt lõi này.

Chuẩn bị môi trường và công cụ

Để bắt đầu, bạn cần một vài thư viện quen thuộc trong hệ sinh thái Python:

• PyTorch — xây dựng và huấn luyện mạng nơ-ron
• Gymnasium (phiên bản mới của OpenAI Gym) — cung cấp môi trường CartPole
• NumPy và Matplotlib — xử lý số liệu và trực quan hóa
• tqdm — thanh tiến trình cho quá trình thu thập dữ liệu

CartPole là lựa chọn lý tưởng cho người mới vì trạng thái chỉ có 4 chiều (vị trí xe, vận tốc xe, góc thanh, vận tốc góc) và hành động chỉ có 2 lựa chọn (đẩy trái hoặc đẩy phải).

Các bước xây dựng World Model

Bước 1: Thu thập dữ liệu từ môi trường

Trước tiên, bạn cần cho một tác nhân đơn giản (có thể là ngẫu nhiên hoặc dùng heuristic) tương tác với CartPole để thu thập hàng nghìn bước chuyển trạng thái:

• Trạng thái hiện tại s
• Hành động a
• Phần thưởng r
• Trạng thái tiếp theo s'
• Cờ kết thúc done

Đây sẽ là tập dữ liệu huấn luyện cho world model.

Bước 2: Định nghĩa kiến trúc mạng

Một world model tối giản cần hai "đầu ra":

• Mô hình động lực (dynamics model): dự đoán s' từ cặp (s, a)
• Mô hình phần thưởng (reward model): dự đoán r từ cặp (s, a)

Cả hai có thể dùng chung một mạng MLP với các lớp ẩn đủ lớn, sau đó tách thành hai nhánh đầu ra. Hàm mất mát kết hợp sai số bình phương trung bình (MSE) cho cả hai đầu ra.

Bước 3: Huấn luyện và kiểm tra

Chia dữ liệu thành tập huấn luyện và tập kiểm tra. Điểm mấu chốt là bạn phải đánh giá trên dữ liệu chưa từng thấy, vì world model rất dễ bị quá khớp (overfitting) — nó có thể thuộc lòng các bước chuyển nhưng thất bại thảm hại khi gặp tình huống mới.

Cho mô hình "mơ mộng" qua CartPole

Sau khi huấn luyện, đây là phần thú vị nhất: bạn không cần môi trường thật nữa. Thay vào đó:

1. Bắt đầu từ một trạng thái khởi tạo thực tế
2. Dùng một chính sách (policy) để chọn hành động
3. Đưa trạng thái và hành động vào world model để dự đoán trạng thái kế tiếp
4. Lặp lại — mô hình tự "tưởng tượng" ra cả một ván CartPole

Kết quả thường rất ấn tượng trong vài chục bước đầu. Nhưng rồi sai số tích lũy bắt đầu bào mòn độ chính xác. Thanh sẽ nghiêng dần theo cách mà vật lý thật không hề cho phép, và cuối cùng ảo giác sụp đổ hoàn toàn.

Đo lường thời điểm ảo giác sụp đổ

Đây là phần mà nhiều hướng dẫn bỏ qua, nhưng lại quan trọng nhất về mặt khoa học. Bạn cần một thước đo định lượng cho "ảo giác" này:

• Sai số dự đoán theo từng bước (one-step error): sai số khi dự đoán chỉ một bước — luôn nhỏ
• Sai số cuốn chiếu (rollout error): sai số khi để mô hình tự chạy nhiều bước liên tiếp — tăng theo hàm mũ
• Chỉ số sụp đổ: thời điểm sai số vượt quá một ngưỡng xác định, hoặc khi mô hình dự đoán "done" sai lệch so với môi trường thật

Cách trực quan hóa hiệu quả nhất là vẽ đồ thị sai số theo số bước rollout. Bạn sẽ thấy một đường cong tăng vọt đặc trưng — dấu hiệu cho thấy mô hình đã đi quá xa khỏi vùng dữ liệu mà nó từng học.

Một world model tốt không phải là mô hình không bao giờ sai, mà là mô hình biết mình đang sai — hoặc ít nhất, ta đo được chính xác nó sai từ lúc nào.

Bài học cho người mới

Sau khi hoàn thành dự án nhỏ này, bạn sẽ nắm được những nguyên lý cốt lõi:

• World model không phải phép thuật — nó chỉ là mạng nơ-ron học các bước chuyển trạng thái
• Sai số tích lũy là kẻ thù số một — mô hình càng mơ xa, ảo giác càng mong manh
• Đánh giá quan trọng không kém huấn luyện — đo lường chính xác giúp bạn hiểu giới hạn thật của mô hình

Ý nghĩa với cộng đồng công nghệ Việt Nam

Với các bạn sinh viên và kỹ sư AI tại Việt Nam, world model là một hướng đi đầy tiềm năng. Nó không đòi hỏi GPU khổng lồ như huấn luyện các mô hình ngôn ngữ lớn, nhưng lại rèn luyện tư duy hệ thống rất tốt. Từ CartPole, bạn có thể mở rộng lên các môi trường phức tạp hơn như Atari, mô phỏng robot, hay thậm chí các bài toán tối ưu trong sản xuất.

Hãy bắt đầu nhỏ, đo lường cẩn thận, và để trí tưởng tượng của mô hình — cũng như của chính bạn — bay xa.