Vào năm 2018, ý tưởng về học tăng cường trong bối cảnh mô hình thế giới thần kinh được giới thiệu lần đầu tiên, và không lâu sau đó, nguyên tắc cơ bản này được áp dụng cho các mô hình thế giới. Một số mô hình nổi bật thực hiện học tăng cường là khuôn khổ Dreamer, đã giới thiệu học tăng cường từ không gian latent của một mô hình không gian trạng thái tái phát. DreamerV2 đã chứng minh rằng việc sử dụng latent rời rạc có thể dẫn đến giảm sai sót tích lũy, và khuôn khổ DreamerV3 đã có thể đạt được hiệu suất giống con người trên một loạt nhiệm vụ trên các lĩnh vực khác nhau với siêu tham số cố định.

Hơn nữa, có thể vẽ ra những điểm tương đồng giữa các mô hình tạo ảnh và mô hình thế giới, cho thấy rằng tiến bộ trong các mô hình tầm nhìn tạo ra có thể được nhân rộng để mang lại lợi ích cho các mô hình thế giới. Kể từ khi sử dụng các bộ chuyển đổi trong các khuôn khổ xử lý ngôn ngữ tự nhiên trở nên phổ biến, các khuôn khổ DALL-E và VQGAN đã xuất hiện. Các khuôn khổ này đã thực hiện các bộ tự động hóa rời rạc để chuyển đổi ảnh thành các token rời rạc, và đã xây dựng các mô hình tạo ảnh từ văn bản rất mạnh mẽ và hiệu quả bằng cách tận dụng khả năng mô hình hóa trình tự của các bộ chuyển đổi tự hồi quy. Đồng thời, các mô hình khuếch tán đã thu hút được sự chú ý, và ngày nay, các mô hình khuếch tán đã thiết lập mình như một khuôn khổ thống trị cho việc tạo ảnh độ phân giải cao. Nhờ vào khả năng của các mô hình khuếch tán và học tăng cường, các nỗ lực đang được thực hiện để kết hợp hai cách tiếp cận này, với mục tiêu tận dụng sự linh hoạt của các mô hình khuếch tán như các mô hình quỹ đạo, mô hình phần thưởng, nhà hoạch định và chính sách cho tăng cường dữ liệu ngoại tuyến.

Các mô hình thế giới cung cấp một phương pháp đầy hứa hẹn cho việc đào tạo các tác nhân học tăng cường một cách an toàn và hiệu quả. Thông thường, các mô hình này sử dụng các chuỗi biến latents rời rạc để mô phỏng động lực học môi trường. Tuy nhiên, sự nén này có thể bỏ qua các chi tiết trực quan quan trọng cho việc học tăng cường. Đồng thời, các mô hình khuếch tán đã trở nên phổ biến cho việc tạo ảnh, thách thức các phương pháp truyền thống sử dụng latents rời rạc. Lấy cảm hứng từ sự thay đổi này, trong bài viết này, chúng tôi sẽ nói về DIAMOND (DIffusion As a Model Of eNvironment Dreams), một tác nhân học tăng cường được đào tạo trong một mô hình thế giới khuếch tán. Chúng tôi sẽ khám phá các lựa chọn thiết kế cần thiết để làm cho khuếch tán phù hợp với mô hình thế giới và chứng minh rằng các chi tiết trực quan được cải thiện dẫn đến hiệu suất tác nhân tốt hơn. DIAMOND thiết lập một điểm chuẩn mới trên thử nghiệm Atari 100k cạnh tranh, đạt được điểm trung bình được chuẩn hóa của con người là 1,46, cao nhất cho các tác nhân được đào tạo hoàn toàn trong một mô hình thế giới.

DIAMOND: DIffusion As a Model Of eNvironment Dreams

Các mô hình thế giới hoặc các mô hình tạo ra môi trường đang nổi lên như một trong những thành phần quan trọng nhất cho các tác nhân tạo ra để lập kế hoạch và suy luận về môi trường của chúng. Mặc dù việc sử dụng học tăng cường đã đạt được thành công đáng kể trong những năm gần đây, các mô hình thực hiện học tăng cường được biết đến là không hiệu quả về mẫu, điều này hạn chế đáng kể ứng dụng thực tế của chúng. Mặt khác, các mô hình thế giới đã chứng minh khả năng đào tạo các tác nhân học tăng cường hiệu quả trên các môi trường đa dạng với hiệu quả mẫu được cải thiện đáng kể, cho phép mô hình học từ kinh nghiệm thực tế. Các khuôn khổ mô hình thế giới gần đây thường mô hình hóa động lực học môi trường như một chuỗi biến latents rời rạc, với mô hình rời rạc hóa không gian latents để tránh sai sót tích lũy trên các đường chân trời thời gian đa bước. Mặc dù cách tiếp cận này có thể mang lại kết quả đáng kể, nhưng nó cũng liên quan đến mất thông tin, dẫn đến mất chất lượng tái tạo và mất tính tổng quát. Mất thông tin có thể trở thành một trở ngại đáng kể cho các kịch bản thực tế đòi hỏi thông tin phải được xác định rõ ràng, như đào tạo xe tự hành. Trong những nhiệm vụ như vậy, các thay đổi nhỏ hoặc chi tiết trong đầu vào trực quan như màu sắc của đèn giao thông hoặc chỉ dẫn quay xe của xe phía trước có thể thay đổi chính sách của một tác nhân. Mặc dù việc tăng số lượng latents rời rạc có thể giúp tránh mất thông tin, nhưng nó cũng làm tăng đáng kể chi phí tính toán.

Hơn nữa, trong những năm gần đây, các mô hình khuếch tán đã xuất hiện như một cách tiếp cận thống trị cho các khuôn khổ tạo ảnh chất lượng cao, vì các khuôn khổ được xây dựng trên các mô hình khuếch tán học cách đảo ngược quá trình làm ồn, và cạnh tranh trực tiếp với một số cách tiếp cận đã được thiết lập để mô hình hóa các token rời rạc, do đó cung cấp một giải pháp thay thế hứa hẹn để loại bỏ nhu cầu rời rạc hóa trong mô hình thế giới. Các mô hình khuếch tán được biết đến với khả năng có thể được điều kiện hóa dễ dàng và mô hình hóa linh hoạt các phân phối đa phương thức phức tạp mà không bị sụp đổ chế độ. Những thuộc tính này rất quan trọng cho mô hình thế giới, vì việc điều kiện hóa cho phép một mô hình thế giới phản ánh chính xác các hành động của một tác nhân, dẫn đến việc phân bổ tín dụng đáng tin cậy hơn. Hơn nữa, việc mô hình hóa các phân phối đa phương thức cung cấp sự đa dạng lớn hơn của các kịch bản đào tạo cho tác nhân, nâng cao hiệu suất tổng thể của nó.

Xây dựng trên những đặc điểm này, DIAMOND, (DIffusion As a Model Of eNvironment Dreams), một tác nhân học tăng cường được đào tạo trong một mô hình thế giới khuếch tán. Khuôn khổ DIAMOND thực hiện các lựa chọn thiết kế cẩn thận để đảm bảo mô hình thế giới khuếch tán của nó vẫn hiệu quả và ổn định trên các đường chân trời thời gian dài. Khuôn khổ cung cấp một phân tích định tính để chứng minh tầm quan trọng của những lựa chọn thiết kế này. DIAMOND thiết lập một điểm chuẩn mới với điểm trung bình được chuẩn hóa của con người là 1,46 trên thử nghiệm Atari 100k được thiết lập, cao nhất cho các tác nhân được đào tạo hoàn toàn trong một mô hình thế giới. Hoạt động trong không gian ảnh cho phép mô hình thế giới khuếch tán của DIAMOND thay thế môi trường một cách liền mạch, cung cấp thông tin sâu hơn về hành vi của mô hình thế giới và tác nhân. Đặc biệt, hiệu suất được cải thiện trong một số trò chơi được quy cho việc mô hình hóa tốt hơn các chi tiết trực quan quan trọng. Khuôn khổ DIAMOND mô hình hóa môi trường như một POMDP hoặc quy trình quyết định Markov có thể quan sát được một phần với một tập hợp các trạng thái, một tập hợp các hành động rời rạc và một tập hợp các quan sát ảnh. Các hàm chuyển đổi mô tả động lực học môi trường, và hàm phần thưởng ánh xạ các chuyển đổi sang phần thưởng vô hướng. Hàm quan sát mô tả xác suất quan sát, và phát ra các quan sát ảnh, được sử dụng bởi các tác nhân để quan sát môi trường, vì chúng không thể truy cập trực tiếp vào các trạng thái. Mục tiêu chính của cách tiếp cận này là để có được một chính sách ánh xạ quan sát sang hành động với nỗ lực tối đa hóa phần thưởng chiết khấu dự kiến với một yếu tố chiết khấu. Các mô hình thế giới là các mô hình tạo ra môi trường, và các mô hình thế giới có thể được sử dụng để tạo ra các môi trường mô phỏng để đào tạo các tác nhân học tăng cường trong môi trường thực và đào tạo các tác nhân học tăng cường trong môi trường mô hình thế giới. Hình 1 minh họa việc unrolling tưởng tượng của khuôn khổ DIAMOND theo thời gian.

DIAMOND: Phương pháp và Kiến trúc

Tại lõi của nó, các mô hình khuếch tán là một lớp các mô hình tạo ra mà tạo ra một mẫu bằng cách đảo ngược quá trình làm ồn, và lấy cảm hứng mạnh mẽ từ nhiệt động lực học không cân bằng. Khuôn khổ DIAMOND xem xét một quá trình khuếch tán được lập chỉ mục bởi một biến thời gian liên tục với các biên và điều kiện biên tương ứng với một phân phối trước không có cấu trúc có thể tính toán được. Hơn nữa, để có được một mô hình tạo ra, mà ánh xạ từ tiếng ồn sang dữ liệu, khuôn khổ DIAMOND phải đảo ngược quá trình, với quá trình đảo ngược cũng là một quá trình khuếch tán, chạy ngược thời gian. Hơn nữa, tại bất kỳ điểm thời gian nào, không phải là một việc tầm thường để ước tính hàm điểm vì khuôn khổ DIAMOND không có quyền truy cập vào hàm điểm thực, và mô hình vượt qua chướng ngại vật này bằng cách thực hiện mục tiêu khớp điểm, một cách tiếp cận cho phép một khuôn khổ đào tạo một mô hình điểm mà không cần biết hàm điểm cơ bản. Mô hình khuếch tán dựa trên điểm cung cấp một mô hình tạo ra vô điều kiện. Tuy nhiên, một mô hình tạo ra điều kiện của động lực học môi trường là cần thiết để phục vụ như một mô hình thế giới, và để phục vụ mục đích này, khuôn khổ DIAMOND xem xét trường hợp chung của cách tiếp cận POMDP, trong đó khuôn khổ có thể sử dụng các quan sát và hành động trong quá khứ để ước tính và tạo ra quan sát tiếp theo trực tiếp. Như được minh họa trong Hình 1, khuôn khổ DIAMOND sử dụng lịch sử này để điều kiện hóa một mô hình khuếch tán, để ước tính và tạo ra quan sát tiếp theo. Mặc dù khuôn khổ DIAMOND có thể sử dụng bất kỳ bộ giải SDE hoặc ODE nào trong lý thuyết, nhưng có một sự đánh đổi giữa NFE hoặc Số lượng đánh giá hàm và chất lượng mẫu ảnh hưởng đến chi phí suy luận của các mô hình khuếch tán đáng kể.

Xây dựng trên những hiểu biết trên, hãy xem xét việc hiện thực hóa khuôn khổ DIAMOND của một mô hình thế giới dựa trên khuếch tán, bao gồm cả hệ số trôi dạt và khuếch tán tương ứng với một lựa chọn cụ thể của cách tiếp cận khuếch tán. Thay vì chọn DDPM, một ứng cử viên tự nhiên phù hợp cho nhiệm vụ, khuôn khổ DIAMOND xây dựng trên công thức EDM, và xem xét một hạt nhân nhiễu với một hàm thực có giá trị của thời gian khuếch tán gọi là lịch trình tiếng ồn. Khuôn khổ chọn các tiền điều kiện để giữ sự biến đổi của đầu vào và đầu ra cho bất kỳ mức tiếng ồn nào. Việc đào tạo mạng kết hợp tín hiệu và tiếng ồn một cách thích ứng tùy thuộc vào mức độ suy giảm, và khi tiếng ồn thấp, mục tiêu trở thành sự khác biệt giữa tín hiệu sạch và tín hiệu bị làm ồn, tức là tiếng ồn Gaussian được thêm vào. Một cách trực giác, điều này ngăn chặn mục tiêu đào tạo trở nên tầm thường trong chế độ tiếng ồn thấp. Trong thực tế, mục tiêu này có phương sai cao ở các cực của lịch trình tiếng ồn, vì vậy mô hình lấy mẫu mức tiếng ồn từ một phân phối log-normal được chọn một cách thực nghiệm để kết hợp đào tạo xung quanh các khu vực tiếng ồn trung bình. Khuôn khổ DIAMOND sử dụng một thành phần U-Net 2D tiêu chuẩn cho trường vectơ, và giữ một bộ đệm các quan sát và hành động trong quá khứ mà khuôn khổ sử dụng để điều kiện hóa. Khuon khổ DIAMOND sau đó kết hợp các quan sát trong quá khứ này với quan sát tiếng ồn tiếp theo và hành động đầu vào thông qua các lớp chuẩn hóa nhóm thích ứng trong các khối dư của U-Net.

DIAMOND: Thử nghiệm và Kết quả

Để đánh giá toàn diện, khuôn khổ DIAMOND chọn thử nghiệm Atari 100k. Thử nghiệm Atari 100k bao gồm 26 trò chơi được thiết kế để kiểm tra một loạt các khả năng của tác nhân. Trong mỗi trò chơi, một tác nhân bị giới hạn ở 100k hành động trong môi trường, tương đương với khoảng 2 giờ chơi trò chơi của con người, để học trò chơi trước khi đánh giá. Để so sánh, các tác nhân Atari không bị giới hạn thường đào tạo trong 50 triệu bước, đại diện cho một sự tăng 500 lần về kinh nghiệm. Chúng tôi đã đào tạo DIAMOND từ đầu sử dụng 5 hạt giống ngẫu nhiên cho mỗi trò chơi. Mỗi lần chạy đào tạo yêu cầu khoảng 12GB VRAM và mất khoảng 2,9 ngày trên một card đồ họa Nvidia RTX 4090, tương đương với 1,03 năm GPU tổng cộng. Bảng sau cung cấp điểm số cho tất cả các trò chơi, điểm trung bình và điểm trung vị của điểm số được chuẩn hóa của con người.

Theo sau những hạn chế của ước tính điểm, khuôn khổ DIAMOND cung cấp sự tự tin bootstrap phân tầng trong điểm trung bình và điểm trung vị của điểm số được chuẩn hóa của con người cùng với các hồ sơ hiệu suất và các chỉ số bổ sung, như được tóm tắt trong hình sau.

Kết quả cho thấy DIAMOND hoạt động đặc biệt tốt trên toàn bộ thử nghiệm, vượt qua người chơi con người trong 11 trò chơi và đạt được điểm trung bình được chuẩn hóa của con người là 1,46, thiết lập một kỷ lục mới cho các tác nhân được đào tạo hoàn toàn trong một mô hình thế giới. Hơn nữa, điểm trung vị của DIAMOND tương đương với STORM và vượt qua tất cả các baseline khác. DIAMOND vượt trội trong các môi trường mà việc bắt捉 các chi tiết nhỏ là quan trọng, chẳng hạn như Asterix, Breakout và RoadRunner. Hơn nữa, như đã thảo luận trước đó, khuôn khổ DIAMOND có tính linh hoạt để thực hiện bất kỳ mô hình khuếch tán nào trong đường ống của nó, mặc dù nó chọn cách tiếp cận EDM, nó sẽ là một lựa chọn tự nhiên để chọn mô hình DDPM vì nó đã được thực hiện trong nhiều ứng dụng tạo ảnh. Để so sánh cách tiếp cận EDM với việc thực hiện DDPM, khuôn khổ DIAMOND đào tạo cả hai biến thể với cùng một kiến trúc mạng trên cùng một tập dữ liệu tĩnh được chia sẻ với hơn 100k khung hình được thu thập với một chính sách chuyên gia. Số lượng bước làm ồn trực tiếp liên quan đến chi phí suy luận của mô hình thế giới, và do đó, ít bước hơn sẽ giảm chi phí đào tạo một tác nhân trên các đường đi tưởng tượng. Để đảm bảo mô hình thế giới của chúng tôi vẫn tính toán so với các baseline khác, chẳng hạn như IRIS yêu cầu 16 NFE mỗi bước thời gian, chúng tôi nhằm mục đích sử dụng không quá vài chục bước làm ồn, lý tưởng là ít hơn. Tuy nhiên, việc đặt số bước làm ồn quá thấp có thể làm giảm chất lượng trực quan, dẫn đến sai sót tích lũy. Để đánh giá sự ổn định của các biến thể khuếch tán khác nhau, chúng tôi hiển thị các đường đi tưởng tượng được tạo ra tự hồi quy lên đến t = 1000 bước thời gian trong hình sau, sử dụng các số bước làm ồn khác nhau n ≤ 10.

Chúng tôi quan sát thấy rằng việc sử dụng DDPM (a) trong chế độ này dẫn đến sai sót tích lũy nghiêm trọng, khiến mô hình thế giới nhanh chóng trôi ra khỏi phân phối. Ngược lại, mô hình thế giới khuếch tán dựa trên EDM (b) vẫn ổn định hơn nhiều trên các đường chân trời thời gian dài, thậm chí với một bước làm ồn. Các đường đi tưởng tượng với các mô hình thế giới khuếch tán dựa trên DDPM (trái) và EDM (phải) được hiển thị. Quan sát ban đầu tại t = 0 là giống nhau cho cả hai, và mỗi hàng tương ứng với một số bước làm ồn giảm. Chúng tôi quan sát thấy rằng việc tạo ra dựa trên DDPM gặp phải sai sót tích lũy, với số bước làm ồn nhỏ hơn dẫn đến tích lũy lỗi nhanh hơn. Ngược lại, mô hình thế giới khuếch tán của DIAMOND vẫn ổn định hơn nhiều, thậm chí với n = 1. Dự đoán bước đơn tối ưu là kỳ vọng trên các tái tạo có thể cho một đầu vào tiếng ồn nhất định, có thể nằm ngoài phân phối nếu phân phối hậu kỳ là đa phương thức. Mặc dù một số trò chơi, như Breakout, có chuyển đổi xác định có thể được mô hình hóa chính xác với một bước làm ồn, các trò chơi khác thể hiện sự không thể quan sát được một phần, dẫn đến phân phối quan sát đa phương thức. Trong những trường hợp này, một bộ giải lặp là cần thiết để hướng dẫn quá trình lấy mẫu tới một chế độ cụ thể, như được minh họa trong trò chơi Boxing trong hình sau. Do đó, khuôn khổ DIAMOND đặt n = 3 trong tất cả các thí nghiệm của chúng tôi.

Hình trên so sánh việc lấy mẫu bước đơn (hàng trên) và lấy mẫu đa bước (hàng dưới) trong Boxing. Các chuyển động của người chơi đen là không thể đoán trước, khiến việc lấy mẫu bước đơn nội suy giữa các kết quả có thể, dẫn đến dự đoán mờ. Ngược lại, việc lấy mẫu đa bước tạo ra một hình ảnh rõ ràng bằng cách hướng dẫn việc tạo ra tới một chế độ cụ thể. Đặc biệt, vì chính sách kiểm soát người chơi trắng, các hành động của anh ta được biết đến bởi mô hình thế giới, loại bỏ sự không chắc chắn. Do đó, cả việc lấy mẫu bước đơn và đa bước đều dự đoán chính xác vị trí của người chơi trắng.

Trong hình trên, các đường đi tưởng tượng bởi DIAMOND thường hiển thị chất lượng trực quan cao hơn và trung thực hơn với môi trường thực so với những đường đi tưởng tượng bởi IRIS. Các đường đi được tạo ra bởi IRIS chứa các không nhất quán trực quan giữa các khung hình (được đánh dấu bằng hộp trắng), chẳng hạn như kẻ thù được hiển thị như phần thưởng và ngược lại. Mặc dù những không nhất quán này có thể chỉ ảnh hưởng đến một vài pixel, nhưng chúng có thể ảnh hưởng đáng kể đến việc học một chính sách tối ưu. Hình hiển thị các khung hình liên tiếp được tưởng tượng với IRIS (trái) và DIAMOND (phải). Các hộp trắng đánh dấu các không nhất quán giữa các khung hình, chỉ xuất hiện trong các đường đi được tạo ra với IRIS. Trong Asterix (hàng trên), một kẻ thù (cam) trở thành phần thưởng (đỏ) trong khung hình thứ hai, sau đó trở lại thành kẻ thù trong khung hình thứ ba và lại trở thành phần thưởng trong khung hình thứ tư. Trong Breakout (hàng giữa), các gạch và điểm số không nhất quán giữa các khung hình. Trong Road Runner (hàng dưới), các phần thưởng (điểm xanh nhỏ trên đường) được hiển thị không nhất quán giữa các khung hình. Những không nhất quán này không xảy ra với DIAMOND. Trong Breakout, điểm số được cập nhật một cách đáng tin cậy bằng +7 khi một gạch đỏ bị phá vỡ.

Kết luận

Trong bài viết này, chúng tôi đã nói về DIAMOND, một tác nhân học tăng cường được đào tạo trong một mô hình thế giới khuếch tán. Khuôn khổ DIAMOND thực hiện các lựa chọn thiết kế cẩn thận để đảm bảo mô hình thế giới khuếch tán của nó vẫn hiệu quả và ổn định trên các đường chân trời thời gian dài. Khuon khổ cung cấp một phân tích định tính để chứng minh tầm quan trọng của những lựa chọn thiết kế này. DIAMOND thiết lập một điểm chuẩn mới với điểm trung bình được chuẩn hóa của con người là 1,46 trên thử nghiệm Atari 100k được thiết lập, cao nhất cho các tác nhân được đào tạo hoàn toàn trong một mô hình thế giới. Hoạt động trong không gian ảnh cho phép mô hình thế giới khuếch tán của DIAMOND thay thế môi trường một cách liền mạch, cung cấp thông tin sâu hơn về hành vi của mô hình thế giới và tác nhân. Đặc biệt, hiệu suất được cải thiện trong một số trò chơi được quy cho việc mô hình hóa tốt hơn các chi tiết trực quan quan trọng. Khuon khổ DIAMOND mô hình hóa môi trường như một POMDP hoặc quy trình quyết định Markov có thể quan sát được một phần với một tập hợp các trạng thái, một tập hợp các hành động rời rạc và một tập hợp các quan sát ảnh. Các hàm chuyển đổi mô tả động lực học môi trường, và hàm phần thưởng ánh xạ các chuyển đổi sang phần thưởng vô hướng.