訓(xùn)練提速17%，第四范式開源強化學(xué)習(xí)研究框架，支持單、多智能體訓(xùn)練

機器之心專欄

作者：第四范式強化學(xué)習(xí)團隊

強化學(xué)習(xí)研究框架 OpenRL 是基于 PyTorch 開發(fā)的，已經(jīng)在 GitHub 上開源。

OpenRL 是由第四范式強化學(xué)習(xí)團隊開發(fā)的基于 PyTorch 的強化學(xué)習(xí)研究框架，支持單智能體、多智能體、自然語言等多種任務(wù)的訓(xùn)練。OpenRL 基于 PyTorch 進行開發(fā)，目標是為強化學(xué)習(xí)研究社區(qū)提供一個簡單易用、靈活高效、可持續(xù)擴展的平臺。目前，OpenRL 支持的特性包括：

簡單易用且支持單智能體、多智能體訓(xùn)練的通用接口

支持自然語言任務(wù)（如對話任務(wù)）的強化學(xué)習(xí)訓(xùn)練

支持從 Hugging Face 上導(dǎo)入模型和數(shù)據(jù)

支持 LSTM，GRU，Transformer 等模型

支持多種訓(xùn)練加速，例如：自動混合精度訓(xùn)練，半精度策略網(wǎng)絡(luò)收集數(shù)據(jù)等

支持用戶自定義訓(xùn)練模型、獎勵模型、訓(xùn)練數(shù)據(jù)以及環(huán)境

支持 gymnasium 環(huán)境

支持字典觀測空間

支持 wandb，tensorboardX 等主流訓(xùn)練可視化工具

支持環(huán)境的串行和并行訓(xùn)練，同時保證兩種模式下的訓(xùn)練效果一致

中英文文檔

提供單元測試和代碼覆蓋測試

符合 Black Code Style 和類型檢查

目前，OpenRL 已經(jīng)在 GitHub 開源：

項目地址：https://github.com/OpenRL-Lab/openrl

OpenRL 初體驗

OpenRL 目前可以通過 pip 進行安裝：

pip install openrl

也可以通過 conda 安裝：

conda install -c openrl openrl

OpenRL 為強化學(xué)習(xí)入門用戶提供了簡單易用的接口， 下面是一個使用 PPO 算法訓(xùn)練 CartPole 環(huán)境的例子：

# train_ppo.pyfrom openrl.envs.common import makefrom openrl.modules.common import PPONet as Netfrom openrl.runners.common import PPOAgent as Agentenv = make ("CartPole-v1", env_num=9) # 創(chuàng)建環(huán)境，并設(shè)置環(huán)境并行數(shù)為 9net = Net (env) # 創(chuàng)建神經(jīng)網(wǎng)絡(luò)agent = Agent (net) # 初始化智能體agent.train (total_time_steps=20000) # 開始訓(xùn)練，并設(shè)置環(huán)境運行總步數(shù)為 20000

使用 OpenRL 訓(xùn)練智能體只需要簡單的四步：創(chuàng)建環(huán)境 => 初始化模型 => 初始化智能體 => 開始訓(xùn)練！

在普通筆記本電腦上執(zhí)行以上代碼，只需要幾秒鐘，便可以完成該智能體的訓(xùn)練：

此外，對于多智能體、自然語言等任務(wù)的訓(xùn)練，OpenRL 也提供了同樣簡單易用的接口。例如，對于多智能體任務(wù)中的 MPE 環(huán)境，OpenRL 也只需要調(diào)用幾行代碼便可以完成訓(xùn)練：

# train_ppo.pyfrom openrl.envs.common import makefrom openrl.modules.common import PPONet as Netfrom openrl.runners.common import PPOAgent as Agentdef train (): # 創(chuàng)建 MPE 環(huán)境，使用異步環(huán)境，即每個智能體獨立運行 env = make ( "simple_spread", env_num=100, asynchronous=True, ) # 創(chuàng)建 神經(jīng)網(wǎng)絡(luò)，使用 GPU 進行訓(xùn)練 net = Net (env, device="cuda") agent = Agent (net) # 初始化訓(xùn)練器 # 開始訓(xùn)練 agent.train (total_time_steps=5000000) # 保存訓(xùn)練完成的智能體 agent.save ("./ppo_agent/")if __name__ == "__main__": train ()

下圖展示了通過 OpenRL 訓(xùn)練前后智能體的表現(xiàn)：

加載配置文件

此外，OpenRL 還同時支持從命令行和配置文件對訓(xùn)練參數(shù)進行修改。比如，用戶可以通過執(zhí)行 python train_ppo.py --lr 5e-4 來快速修改訓(xùn)練時候的學(xué)習(xí)率。

當配置參數(shù)非常多的時候，OpenRL 還支持用戶編寫自己的配置文件來修改訓(xùn)練參數(shù)。例如，用戶可以自行創(chuàng)建以下配置文件 (mpe_ppo.yaml)，并修改其中的參數(shù)：

# mpe_ppo.yamlseed: 0 # 設(shè)置 seed，保證每次實驗結(jié)果一致lr: 7e-4 # 設(shè)置學(xué)習(xí)率episode_length: 25 # 設(shè)置每個 episode 的長度use_recurrent_policy: true # 設(shè)置是否使用 RNNuse_joint_action_loss: true # 設(shè)置是否使用 JRPO 算法use_valuenorm: true # 設(shè)置是否使用 value normalization

最后，用戶只需要在執(zhí)行程序的時候指定該配置文件即可：

python train_ppo.py --config mpe_ppo.yaml

訓(xùn)練與測試可視化

此外，通過 OpenRL，用戶還可以方便地使用 wandb 來可視化訓(xùn)練過程：

OpenRL 還提供了各種環(huán)境可視化的接口，方便用戶對并行環(huán)境進行可視化。用戶可以在創(chuàng)建并行環(huán)境的時候設(shè)置環(huán)境的渲染模式為 "group_human"，便可以同時對多個并行環(huán)境進行可視化：

env = make ("simple_spread", env_num=9, render_mode="group_human")

此外，用戶還可以通過引入 GIFWrapper 來把環(huán)境運行過程保存為 gif 動畫：

from openrl.envs.wrappers import GIFWrapperenv = GIFWrapper (env, "test_simple_spread.gif")

智能體的保存和加載

OpenRL 提供 agent.save () 和 agent.load () 接口來保存和加載訓(xùn)練好的智能體，并通過 agent.act () 接口來獲取測試時的智能體動作：

# test_ppo.pyfrom openrl.envs.common import makefrom openrl.modules.common import PPONet as Netfrom openrl.runners.common import PPOAgent as Agentfrom openrl.envs.wrappers import GIFWrapper # 用于生成 gifdef test (): # 創(chuàng)建 MPE 環(huán)境 env = make ( "simple_spread", env_num=4) # 使用 GIFWrapper，用于生成 gif env = GIFWrapper (env, "test_simple_spread.gif") agent = Agent (Net (env)) # 創(chuàng)建 智能體 # 保存智能體 agent.save ("./ppo_agent/") # 加載智能體 agent.load ('./ppo_agent/') # 開始測試 obs, _ = env.reset () while True: # 智能體根據(jù) observation 預(yù)測下一個動作 action, _ = agent.act (obs) obs, r, done, info = env.step (action) if done.any (): break env.close ()if __name__ == "__main__": test ()

執(zhí)行該測試代碼，便可以在同級目錄下找到保存好的環(huán)境運行動畫文件 (test_simple_spread.gif)：

訓(xùn)練自然語言對話任務(wù)

最近的研究表明，強化學(xué)習(xí)也可以用于訓(xùn)練語言模型， 并且能顯著提升模型的性能。目前，OpenRL 已經(jīng)支持自然語言對話任務(wù)的強化學(xué)習(xí)訓(xùn)練。OpenRL 通過模塊化設(shè)計，支持用戶加載自己的數(shù)據(jù)集 ，自定義訓(xùn)練模型，自定義獎勵模型，自定義 wandb 信息輸出以及一鍵開啟混合精度訓(xùn)練等。

對于對話任務(wù)訓(xùn)練，OpenRL 提供了同樣簡單易用的訓(xùn)練接口：

# train_ppo.pyfrom openrl.envs.common import makefrom openrl.modules.common import PPONet as Netfrom openrl.runners.common import PPOAgent as Agentfrom openrl.configs.config import create_config_parserdef train (): # 添加讀取配置文件的代碼 cfg_parser = create_config_parser () cfg = cfg_parser.parse_args () # 創(chuàng)建 NLP 環(huán)境 env = make ("daily_dialog",env_num=2,asynchronous=True,cfg=cfg,) net = Net (env, cfg=cfg, device="cuda") agent = Agent (net) agent.train (total_time_steps=5000000)if __name__ == "__main__": train ()

可以看出，OpenRL 訓(xùn)練對話任務(wù)和其他強化學(xué)習(xí)任務(wù)一樣，都是通過創(chuàng)建交互環(huán)境的方式進行訓(xùn)練。

加載自定義數(shù)據(jù)集

訓(xùn)練對話任務(wù)，需要對話數(shù)據(jù)集。這里我們可以使用 Hugging Face 上的公開數(shù)據(jù)集（用戶可以替換成自己的數(shù)據(jù)集）。加載數(shù)據(jù)集，只需要在配置文件中傳入數(shù)據(jù)集的名稱或者路徑即可：

# nlp_ppo.yamldata_path: daily_dialog # 數(shù)據(jù)集路徑env: # 環(huán)境所用到的參數(shù) args: {'tokenizer_path': 'gpt2'} # 讀取 tokenizer 的路徑seed: 0 # 設(shè)置 seed，保證每次實驗結(jié)果一致lr: 1e-6 # 設(shè)置 policy 模型的學(xué)習(xí)率critic_lr: 1e-6 # 設(shè)置 critic 模型的學(xué)習(xí)率episode_length: 20 # 設(shè)置每個 episode 的長度use_recurrent_policy: true

上述配置文件中的 data_path 可以設(shè)置為 Hugging Face 數(shù)據(jù)集名稱或者本地數(shù)據(jù)集路徑。此外，環(huán)境參數(shù)中的 tokenizer_path 用于指定加載文字編碼器的 Hugging Face 名稱或者本地路徑。

自定義訓(xùn)練模型

在 OpenRL 中，我們可以使用 Hugging Face 上的模型來進行訓(xùn)練。為了加載 Hugging Face 上的模型，我們首先需要在配置文件 nlp_ppo.yaml 中添加以下內(nèi)容：

# nlp_ppo.yaml# 預(yù)訓(xùn)練模型路徑model_path: rajkumarrrk/gpt2-fine-tuned-on-daily-dialog use_share_model: true # 策略網(wǎng)絡(luò)和價值網(wǎng)絡(luò)是否共享模型ppo_epoch: 5 # ppo 訓(xùn)練迭代次數(shù)

data_path: daily_dialog # 數(shù)據(jù)集名稱或者路徑env: # 環(huán)境所用到的參數(shù) args: {'tokenizer_path': 'gpt2'} # 讀取 tokenizer 的路徑lr: 1e-6 # 設(shè)置 policy 模型的學(xué)習(xí)率critic_lr: 1e-6 # 設(shè)置 critic 模型的學(xué)習(xí)率episode_length: 128 # 設(shè)置每個 episode 的長度num_mini_batch: 20

然后在 train_ppo.py 中添加以下代碼：

# train_ppo.pyfrom openrl.envs.common import makefrom openrl.modules.common import PPONet as Netfrom openrl.runners.common import PPOAgent as Agentfrom openrl.configs.config import create_config_parserfrom openrl.modules.networks.policy_value_network_gpt import ( PolicyValueNetworkGPT as PolicyValueNetwork,)def train (): # 添加讀取配置文件的代碼 cfg_parser = create_config_parser () cfg = cfg_parser.parse_args () # 創(chuàng)建 NLP 環(huán)境 env = make ("daily_dialog",env_num=2,asynchronous=True,cfg=cfg,) # 創(chuàng)建自定義神經(jīng)網(wǎng)絡(luò) model_dict = {"model": PolicyValueNetwork} net = Net (env, cfg=cfg, model_dict=model_dict) # 創(chuàng)建訓(xùn)練智能體 agent = Agent (net) agent.train (total_time_steps=5000000)if __name__ == "__main__": train ()

通過以上簡單幾行的修改，用戶便可以使用 Hugging Face 上的預(yù)訓(xùn)練模型進行訓(xùn)練。如果用戶希望分別自定義策略網(wǎng)絡(luò)和價值網(wǎng)絡(luò)，可以寫好 CustomPolicyNetwork 以及 CustomValueNetwork 后通過以下方式從外部傳入訓(xùn)練網(wǎng)絡(luò)：

model_dict = { "policy": CustomPolicyNetwork, "critic": CustomValueNetwork,}net = Net (env, model_dict=model_dict)

自定義獎勵模型

通常，自然語言任務(wù)的數(shù)據(jù)集中并不包含獎勵信息。因此，如果需要使用強化學(xué)習(xí)來訓(xùn)練自然語言任務(wù)，就需要使用額外的獎勵模型來生成獎勵。在該對話任務(wù)中，我們可以使用一個復(fù)合的獎勵模型，它包含以下三個部分：

●意圖獎勵：即當智能體生成的語句和期望的意圖接近時，智能體便可以獲得更高的獎勵。

●METEOR 指標獎勵：METEOR 是一個用于評估文本生成質(zhì)量的指標，它可以用來衡量生成的語句和期望的語句的相似程度。我們把這個指標作為獎勵反饋給智能體，以達到優(yōu)化生成的語句的效果。

●KL 散度獎勵：該獎勵用來限制智能體生成的文本偏離預(yù)訓(xùn)練模型的程度，防止出現(xiàn) reward hacking 的問題。

我們最終的獎勵為以上三個獎勵的加權(quán)和，其中 KL 散度獎勵的系數(shù)是隨著 KL 散度的大小動態(tài)變化的。想在 OpenRL 中使用該獎勵模型，用戶無需修改訓(xùn)練代碼，只需要在 nlp_ppo.yaml 文件中添加 reward_class 參數(shù)即可：

# nlp_ppo.yamlreward_class: id: NLPReward # 獎勵模型名稱 args: { # 用于意圖判斷的模型的名稱或路徑 "intent_model": rajkumarrrk/roberta-daily-dialog-intent-classifier, # 用于計算 KL 散度的預(yù)訓(xùn)練模型的名稱或路徑 "ref_model": roberta-base, # 用于意圖判斷的 tokenizer 的名稱或路徑 }

OpenRL 支持用戶使用自定義的獎勵模型。首先，用戶需要編寫自定義獎勵模型 (需要繼承 BaseReward 類)。接著，用戶需要注冊自定義的獎勵模型，即在 train_ppo.py 添加以下代碼：

# train_ppo.pyfrom openrl.rewards.nlp_reward import CustomRewardfrom openrl.rewards import RewardFactoryRewardFactory.register ("CustomReward", CustomReward)

最后，用戶只需要在配置文件中填寫自定義的獎勵模型即可：

reward_class: id: "CustomReward" # 自定義獎勵模型名稱 args: {} # 用戶自定義獎勵函數(shù)可能用到的參數(shù)

自定義訓(xùn)練過程信息輸出

OpenRL 還支持用戶自定義 wandb 和 tensorboard 的輸出內(nèi)容。例如，在該任務(wù)的訓(xùn)練過程中，我們還需要輸出各種類型獎勵的信息和 KL 散度系數(shù)的信息， 用戶可以在 nlp_ppo.yaml 文件中加入 vec_info_class 參數(shù)來實現(xiàn):

# nlp_ppo.yamlvec_info_class: id: "NLPVecInfo" # 調(diào)用 NLPVecInfo 類以打印 NLP 任務(wù)中獎勵函數(shù)的信息# 設(shè)置 wandb 信息wandb_entity: openrl # 這里用于指定 wandb 團隊名稱，請把 openrl 替換為你自己的團隊名稱experiment_name: train_nlp # 這里用于指定實驗名稱run_dir: ./run_results/ # 這里用于指定實驗數(shù)據(jù)保存的路徑log_interval: 1 # 這里用于指定每隔多少個 episode 上傳一次 wandb 數(shù)據(jù)# 自行填寫其他參數(shù)...

修改完配置文件后，在 train_ppo.py 文件中啟用 wandb:

# train_ppo.pyagent.train (total_time_steps=100000, use_wandb=True)

然后執(zhí)行 python train_ppo.py –config nlp_ppo.yaml，稍后，便可以在 wandb 中看到如下的輸出:

從上圖可以看到，wandb 輸出了各種類型獎勵的信息和 KL 散度系數(shù)的信息。

如果用戶還需要輸出其他信息，還可以參考 NLPVecInfo 類 和 VecInfo 類來實現(xiàn)自己的 CustomVecInfo 類。然后，需要在 train_ppo.py 中注冊自定義的 CustomVecInfo 類:

# train_ppo.py # 注冊自定義輸出信息類 VecInfoFactory.register ("CustomVecInfo", CustomVecInfo)

最后，只需要在 nlp_ppo.yaml 中填寫 CustomVecInfo 類即可啟用：

# nlp_ppo.yamlvec_info_class: id: "CustomVecInfo" # 調(diào)用自定義 CustomVecInfo 類以輸出自定義信息

使用混合精度訓(xùn)練加速

OpenRL 還提供了一鍵開啟混合精度訓(xùn)練的功能。用戶只需要在配置文件中加入以下參數(shù)即可：

# nlp_ppo.yamluse_amp: true # 開啟混合精度訓(xùn)練

對比評測

下表格展示了使用 OpenRL 訓(xùn)練該對話任務(wù)的結(jié)果。結(jié)果顯示使用強化學(xué)習(xí)訓(xùn)練后，模型各項指標皆有所提升。另外，從下表可以看出，相較于 RL4LMs ， OpenRL 的訓(xùn)練速度更快（在同樣 3090 顯卡的機器上，速度提升 17% ），最終的性能指標也更好：

最后，對于訓(xùn)練好的智能體，用戶可以方便地通過 agent.chat () 接口進行對話：

# chat.pyfrom openrl.runners.common import ChatAgent as Agentdef chat (): agent = Agent.load ("./ppo_agent", tokenizer="gpt2",) history = [] print ("Welcome to OpenRL!") while True: input_text = input ("> User:") if input_text == "quit": break elif input_text == "reset": history = [] print ("Welcome to OpenRL!") continue response = agent.chat (input_text, history) print (f"> OpenRL Agent: {response}") history.append (input_text) history.append (response)if __name__ == "__main__": chat ()

執(zhí)行 python chat.py ，便可以和訓(xùn)練好的智能體進行對話了：

總結(jié)

OpenRL 框架經(jīng)過了 OpenRL-Lab 的多次迭代并應(yīng)用于學(xué)術(shù)研究和 AI 競賽，目前已經(jīng)成為了一個較為成熟的強化學(xué)習(xí)框架。OpenRL-Lab 團隊將持續(xù)維護和更新 OpenRL，歡迎大家加入我們的開源社區(qū)，一起為強化學(xué)習(xí)的發(fā)展做出貢獻。更多關(guān)于 OpenRL 的信息，可以參考:

OpenRL 官方倉庫：https://github.com/OpenRL-Lab/openrl/

OpenRL 中文文檔：https://openrl-docs.readthedocs.io/zh/latest/

致謝

OpenRL 框架的開發(fā)吸取了其他強化學(xué)習(xí)框架的優(yōu)點：

Stable-baselines3: https://github.com/DLR-RM/stable-baselines3

pytorch-a2c-ppo-acktr-gail:https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail

MAPPO: https://github.com/marlbenchmark/on-policy

Gymnasium: https://github.com/Farama-Foundation/Gymnasium

DI-engine:https://github.com/opendilab/DI-engine/

Tianshou: https://github.com/thu-ml/tianshou

RL4LMs: https://github.com/allenai/RL4LMs

未來工作

目前，OpenRL 還處于持續(xù)開發(fā)和建設(shè)階段，未來 OpenRL 將會開源更多功能：

支持智能體自博弈訓(xùn)練

加入離線強化學(xué)習(xí)、模范學(xué)習(xí)、逆強化學(xué)習(xí)算法

加入更多強化學(xué)習(xí)環(huán)境和算法

集成 Deepspeed 等加速框架

支持多機分布式訓(xùn)練

OpenRL Lab 團隊

OpenRL框架是由OpenRL Lab團隊開發(fā)，該團隊是第四范式公司旗下的強化學(xué)習(xí)研究團隊。第四范式長期致力于強化學(xué)習(xí)的研發(fā)和工業(yè)應(yīng)用。為了促進強化學(xué)習(xí)的產(chǎn)學(xué)研一體化，第四范式成立了OpenRL Lab研究團隊，目標是先進技術(shù)開源和人工智能前沿探索。成立不到一年，OpenRL Lab團隊已經(jīng)在AAMAS發(fā)表過三篇論文，參加谷歌足球游戲 11 vs 11比賽并獲得第三的成績。團隊提出的TiZero智能體，實現(xiàn)了首個從零開始，通過課程學(xué)習(xí)、分布式強化學(xué)習(xí)、自博弈等技術(shù)完成谷歌足球全場游戲智能體的訓(xùn)練：

截止 2022 年 10 月 28 日，Tizero 在及第評測平臺上排名第一：

? THE END

轉(zhuǎn)載請聯(lián)系本公眾號獲得授權(quán)

投稿或?qū)で髨蟮溃?a href="/cdn-cgi/l/email-protection" class="__cf_email__" data-cfemail="bfdcd0d1cbdad1cbffd5d6ced6c5d7d6c7d6d191dcd0d2">[email protected]

原標題：《訓(xùn)練提速17%，第四范式開源強化學(xué)習(xí)研究框架，支持單、多智能體訓(xùn)練》

閱讀原文