引入大語言模型、首個支持國產算力生態,全新開源RL框架RLLTE來了
近年來,強化學習的研究熱度不斷攀升,在智能制造、自動駕駛、大語言模型等多個領域取得了耀眼成績,展示出巨大的研究潛力。然而,強化學習算法高效、可靠的工程實現仍是長期存在的問題。由于其算法結構的復雜性,微小的代碼差異就可能嚴重影響實際性能。
為了解決這一問題,科研和開發者社區先后提出了多個強化學習框架,如強調穩定性和可靠性的 Stable-Baselines3、模塊化設計的 Tianshou 以及單文件實現算法的 CleanRL,為強化學習的學術研究和應用開發做出了積極貢獻。然而,大部分項目的活躍周期較短,未建立合理的長期演進計劃,并且代碼風格迥異,限制了社區的開源協作。它們也沒有構建完整的項目生態,僅專注于模型訓練,而忽略了評估、部署等現實需求。同時,這些項目缺乏完備的測試數據,導致復現成本極高,阻礙了后續研究的進行。
為了解決以上問題,來自香港理工大學、寧波東方理工大學(暫名)、普渡大學和大疆科技的研究者和算法團隊,受到電信中「長期演進技術(LTE)」的啟發,發布了名為 RLLTE 的開源強化學習框架,旨在為促進強化學習研究和應用提供開發組件和工程標準。RLLTE 不僅提供高質量的算法實現,還可作為工具庫用于新算法的開發。
- 論文鏈接:https://arxiv.org/pdf/2309.16382.pdf
- 代碼 / Demo 鏈接:https://github.com/RLE-Foundation/rllte
- 官方網站:https://docs.rllte.dev/
圖 1 RLLTE 框架示意圖
RLLTE 從「探索 - 利用」的角度出發對強化學習算法進行完全解耦,將它們分解為若干最小基元,例如用于處理觀測的編碼器(Encoder),以及用于經驗存儲和采樣的存儲器(Storage)。RLLTE 為每一種基元提供了豐富的模組供開發者選擇,使得開發者可以以「搭積木」的方式進行強化學習算法的構建。RLLTE 框架的主要功能和亮點如下:
- 極致模塊化:RLLTE 的核心設計思想是像 PyTorch 一樣為強化學習算法開發提供標準、便捷、即插即用的開發組件,而并非專注于提供具體的算法實現。因此,RLLTE 實現的算法中每個組件都是可替換的,并且支持用戶使用自定義的模塊。這一解耦過程有助于算法可解釋性的研究和更深層次的改進探索。
- 長期演進:RLLTE 作為一個長期演進的框架,將會持續更新強化學習中的先進算法和工具。為了保持項目的體量和高質量,RLLTE 只更新通用的算法,抑或是在采樣效率或者泛化能力方面做出的改進,并且要求這些算法必須在廣受認可的基線上取得杰出的性能。
- 數據增強:近年來的大量研究將數據增強技巧引入強化學習算法,以實現采樣效率和泛化能力的顯著提升。RLLTE 默認支持數據增強操作,并提供大量觀測(Observations)增強和內在獎勵(Intrinsic Rewards)模組供開發者選擇。
- 豐富的項目生態:RLLTE 同時考慮學術界和工業界的需求,構建了豐富的項目生態。開發者可以在一個框架中實現任務設計、模型訓練、評估以及部署。并且,RLLTE 還嘗試將大語言模型引入該框架,以降低用戶學習成本,加速強化學習的應用構建。
- 完備的基線數據:現有的強化學習框架通常只在有限的任務上進行算法測試,缺乏完備的測試數據(訓練曲線和測試分數等)。由于強化學習訓練的高算力消耗,這是可以理解的,但阻礙了后續的研究進行。為了解決這一問題,RLLTE 依托 Hugging Face 平臺建立了數據倉庫,并對內置算法在廣受認可的基線上進行測試以提供完整的訓練數據。
- 多硬件支持:在當前全球算力需求激增的背景下,RLLTE 被設計為支持多種算力設備以保證靈活性和可拓展性。當前,框架支持使用 NVIDIA GPU 和 HUAWEI NPU 進行訓練,并支持在 NVIDIA TensorRT 以及 HUAWEI CANN 架構下進行推理端部署。RLTLE 也是首個支持國產算力生態的強化學習框架。
圖 2 RLLTE 框架和其他常見開源 RL 框架的功能對比
RLLTE 框架介紹
RLLTE 框架主要包含三個層級:核心層(Core)、應用層(Application)以及工具層(Tool)。
- 核心層從「探索 - 利用」的角度對 RL 算法進行完全解耦,并將其拆分成以下 6 大基元:
圖 3 RLLTE 核心層架構
RLLTE 為每一類基元提供了大量模組供開發者選擇。例如,storage 模塊中提供了 VanillaReplayStorage 和 DictReplayStorage 分別用于存儲普通格式和字典格式的觀測數據。
- 應用層基于核心層的模組提供強化學習算法的實現(rllte.agent)、預訓練(Pre-training)、模型部署(Deployment)以及開發助手(Copilot)。
- 工具層為任務設計(rllte.env)、模型評估(rllte.evaluation)以及基線數據(rllte.hub)提供方便的應用接口。
用戶可以直接調用 RLLTE 實現好的算法進行訓練,例如使用 DrQ-v2 算法解決視覺控制任務:
開始訓練,將會看到以下輸出:
或者,使用內置的算法原型和模塊進行快速算法開發:
使用 RLLTE,開發者只需幾十行代碼即可實現 A2C 等知名強化學習算法。并且,對于實現好的算法,開發者可以更換其內置模塊以進行性能比較和算法改進。例如,比較不同的 Encoder 對同一算法性能的影響:
對于訓練好的模型,開發者可以調用 rllte.evaluation 中的工具進行算法評估和可視化:
在學術研究中,使用 rllte.hub 快速獲取算法的訓練數據和模型參數:
大語言模型融入 RL 框架
RLLTE 是首個嘗試引入大語言模型的強化學習框架,以此來降低開發者的學習成本,以及協助進行強化學習應用的開發。
RLLTE 參考了 LocalGPT 等項目的設計理念,開發了 Copilot,無需額外訓練并且保證用戶數據的絕對隱私。其首先使用 instructor embedding 工具對由項目文檔、教程以及強化學習論文等構成的語料庫進行處理,以建立本地化的向量數據庫,然后使用諸如 Vicuna-7B 的大模型對問題進行理解,并基于該向量數據庫給出答案。用戶可以根據自己的算力情況自由更換基礎模型,未來我們也將進一步豐富該語料庫并添加更多高級功能來實現更加智能的強化學習專屬 Copilot。
作為一個長期演進的強化學習框架,RLLTE 未來將持續跟蹤最新的研究進展并提供高質量的算法實現,以適應開發者變化的需求并對強化學習社區產生積極影響。