天天看點

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

作者:機器之心Pro

機器之心專欄

作者:第四範式強化學習團隊

強化學習研究架構 OpenRL 是基于 PyTorch 開發的,已經在 GitHub 上開源。

OpenRL 是由第四範式強化學習團隊開發的基于 PyTorch 的強化學習研究架構,支援單智能體、多智能體、自然語言等多種任務的訓練。OpenRL 基于 PyTorch 進行開發,目标是為強化學習研究社群提供一個簡單易用、靈活高效、可持續擴充的平台。目前,OpenRL 支援的特性包括:

  • 簡單易用且支援單智能體、多智能體訓練的通用接口
  • 支援自然語言任務(如對話任務)的強化學習訓練
  • 支援從 Hugging Face 上導入模型和資料
  • 支援 LSTM,GRU,Transformer 等模型
  • 支援多種訓練加速,例如:自動混合精度訓練,半精度政策網絡收集資料等
  • 支援使用者自定義訓練模型、獎勵模型、訓練資料以及環境
  • 支援 gymnasium 環境
  • 支援字典觀測空間
  • 支援 wandb,tensorboardX 等主流訓練可視化工具
  • 支援環境的串行和并行訓練,同時保證兩種模式下的訓練效果一緻
  • 中英文文檔
  • 提供單元測試和代碼覆寫測試
  • 符合 Black Code Style 和類型檢查

目前,OpenRL 已經在 GitHub 開源:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

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

OpenRL 初體驗

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

pip install openrl           

也可以通過 conda 安裝:

conda install -c openrl openrl           

OpenRL 為強化學習入門使用者提供了簡單易用的接口, 下面是一個使用 PPO 算法訓練 CartPole 環境的例子:

# train_ppo.py
from openrl.envs.common import make
from openrl.modules.common import PPONet as Net
from openrl.runners.common import PPOAgent as Agent
env = make ("CartPole-v1", env_num=9) # 建立環境,并設定環境并行數為 9
net = Net (env) # 建立神經網絡
agent = Agent (net) # 初始化智能體
agent.train (total_time_steps=20000) # 開始訓練,并設定環境運作總步數為 20000           

使用 OpenRL 訓練智能體隻需要簡單的四步:建立環境 => 初始化模型 => 初始化智能體 => 開始訓練!

在普通筆記本電腦上執行以上代碼,隻需要幾秒鐘,便可以完成該智能體的訓練:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

此外,對于多智能體、自然語言等任務的訓練,OpenRL 也提供了同樣簡單易用的接口。例如,對于多智能體任務中的 MPE 環境,OpenRL 也隻需要調用幾行代碼便可以完成訓練:

# train_ppo.py
from openrl.envs.common import make
from openrl.modules.common import PPONet as Net
from openrl.runners.common import PPOAgent as Agent
def train ():
    # 建立 MPE 環境,使用異步環境,即每個智能體獨立運作
    env = make (
        "simple_spread",
        env_num=100,
        asynchronous=True,
    )
    # 建立 神經網絡,使用 GPU 進行訓練
    net = Net (env, device="cuda")
    agent = Agent (net) # 初始化訓練器
    # 開始訓練
    agent.train (total_time_steps=5000000)
    # 儲存訓練完成的智能體
    agent.save ("./ppo_agent/")
if __name__ == "__main__":
    train ()           

下圖展示了通過 OpenRL 訓練前後智能體的表現:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

加載配置檔案

此外,OpenRL 還同時支援從指令行和配置檔案對訓練參數進行修改。比如,使用者可以通過執行 python train_ppo.py --lr 5e-4 來快速修改訓練時候的學習率。

當配置參數非常多的時候,OpenRL 還支援使用者編寫自己的配置檔案來修改訓練參數。例如,使用者可以自行建立以下配置檔案 (mpe_ppo.yaml),并修改其中的參數:

# mpe_ppo.yaml
seed: 0 # 設定 seed,保證每次實驗結果一緻
lr: 7e-4 # 設定學習率
episode_length: 25 # 設定每個 episode 的長度
use_recurrent_policy: true # 設定是否使用 RNN
use_joint_action_loss: true # 設定是否使用 JRPO 算法
use_valuenorm: true # 設定是否使用 value normalization           

最後,使用者隻需要在執行程式的時候指定該配置檔案即可:

python train_ppo.py --config mpe_ppo.yaml           

訓練與測試可視化

此外,通過 OpenRL,使用者還可以友善地使用 wandb 來可視化訓練過程:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

OpenRL 還提供了各種環境可視化的接口,友善使用者對并行環境進行可視化。使用者可以在建立并行環境的時候設定環境的渲染模式為 "group_human",便可以同時對多個并行環境進行可視化:

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

此外,使用者還可以通過引入 GIFWrapper 來把環境運作過程儲存為 gif 動畫:

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

智能體的儲存和加載

OpenRL 提供 agent.save () 和 agent.load () 接口來儲存和加載訓練好的智能體,并通過 agent.act () 接口來擷取測試時的智能體動作:

# test_ppo.py
from openrl.envs.common import make
from openrl.modules.common import PPONet as Net
from openrl.runners.common import PPOAgent as Agent
from openrl.envs.wrappers import GIFWrapper # 用于生成 gif
def test ():
    # 建立 MPE 環境
    env = make ( "simple_spread", env_num=4)
    # 使用 GIFWrapper,用于生成 gif
    env = GIFWrapper (env, "test_simple_spread.gif")
    agent = Agent (Net (env)) # 建立 智能體
    # 儲存智能體
    agent.save ("./ppo_agent/")    
    # 加載智能體
    agent.load ('./ppo_agent/')
    # 開始測試
    obs, _ = env.reset ()
    while True:
        # 智能體根據 observation 預測下一個動作
        action, _ = agent.act (obs)
        obs, r, done, info = env.step (action)
        if done.any ():
            break
    env.close ()
if __name__ == "__main__":
    test ()           

執行該測試代碼,便可以在同級目錄下找到儲存好的環境運作動畫檔案 (test_simple_spread.gif):

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

訓練自然語言對話任務

最近的研究表明,強化學習也可以用于訓練語言模型, 并且能顯著提升模型的性能。目前,OpenRL 已經支援自然語言對話任務的強化學習訓練。OpenRL 通過子產品化設計,支援使用者加載自己的資料集 ,自定義訓練模型,自定義獎勵模型,自定義 wandb 資訊輸出以及一鍵開啟混合精度訓練等。

對于對話任務訓練,OpenRL 提供了同樣簡單易用的訓練接口:

# train_ppo.py
from openrl.envs.common import make
from openrl.modules.common import PPONet as Net
from openrl.runners.common import PPOAgent as Agent
from openrl.configs.config import create_config_parser
def train ():
    # 添加讀取配置檔案的代碼
    cfg_parser = create_config_parser ()
    cfg = cfg_parser.parse_args ()
    # 建立 NLP 環境
    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 訓練對話任務和其他強化學習任務一樣,都是通過建立互動環境的方式進行訓練。

加載自定義資料集

訓練對話任務,需要對話資料集。這裡我們可以使用 Hugging Face 上的公開資料集(使用者可以替換成自己的資料集)。加載資料集,隻需要在配置檔案中傳入資料集的名稱或者路徑即可:

# nlp_ppo.yaml
data_path: daily_dialog # 資料集路徑
env: # 環境所用到的參數
    args: {'tokenizer_path': 'gpt2'} # 讀取 tokenizer 的路徑
seed: 0 # 設定 seed,保證每次實驗結果一緻
lr: 1e-6 # 設定 policy 模型的學習率
critic_lr: 1e-6 # 設定 critic 模型的學習率
episode_length: 20 # 設定每個 episode 的長度
use_recurrent_policy: true           

上述配置檔案中的 data_path 可以設定為 Hugging Face 資料集名稱或者本地資料集路徑。此外,環境參數中的 tokenizer_path 用于指定加載文字編碼器的 Hugging Face 名稱或者本地路徑。

自定義訓練模型

在 OpenRL 中,我們可以使用 Hugging Face 上的模型來進行訓練。為了加載 Hugging Face 上的模型,我們首先需要在配置檔案 nlp_ppo.yaml 中添加以下内容:

# nlp_ppo.yaml
# 預訓練模型路徑
model_path: rajkumarrrk/gpt2-fine-tuned-on-daily-dialog 
use_share_model: true # 政策網絡和價值網絡是否共享模型
ppo_epoch: 5 # ppo 訓練疊代次數

data_path: daily_dialog # 資料集名稱或者路徑
env: # 環境所用到的參數
    args: {'tokenizer_path': 'gpt2'} # 讀取 tokenizer 的路徑
lr: 1e-6 # 設定 policy 模型的學習率
critic_lr: 1e-6 # 設定 critic 模型的學習率
episode_length: 128 # 設定每個 episode 的長度
num_mini_batch: 20           

然後在 train_ppo.py 中添加以下代碼:

# train_ppo.py
from openrl.envs.common import make
from openrl.modules.common import PPONet as Net
from openrl.runners.common import PPOAgent as Agent
from openrl.configs.config import create_config_parser
from openrl.modules.networks.policy_value_network_gpt import (
    PolicyValueNetworkGPT as PolicyValueNetwork,
)
def train ():
    # 添加讀取配置檔案的代碼
    cfg_parser = create_config_parser ()
    cfg = cfg_parser.parse_args ()
    # 建立 NLP 環境
    env = make ("daily_dialog",env_num=2,asynchronous=True,cfg=cfg,)
    # 建立自定義神經網絡
    model_dict = {"model": PolicyValueNetwork}
    net = Net (env, cfg=cfg, model_dict=model_dict)
    # 建立訓練智能體
    agent = Agent (net)
    agent.train (total_time_steps=5000000)
if __name__ == "__main__":
    train ()           

通過以上簡單幾行的修改,使用者便可以使用 Hugging Face 上的預訓練模型進行訓練。如果使用者希望分别自定義政策網絡和價值網絡,可以寫好 CustomPolicyNetwork 以及 CustomValueNetwork 後通過以下方式從外部傳入訓練網絡:

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

自定義獎勵模型

通常,自然語言任務的資料集中并不包含獎勵資訊。是以,如果需要使用強化學習來訓練自然語言任務,就需要使用額外的獎勵模型來生成獎勵。在該對話任務中,我們可以使用一個複合的獎勵模型,它包含以下三個部分:

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

●METEOR 名額獎勵:METEOR 是一個用于評估文本生成品質的名額,它可以用來衡量生成的語句和期望的語句的相似程度。我們把這個名額作為獎勵回報給智能體,以達到優化生成的語句的效果。

●KL 散度獎勵:該獎勵用來限制智能體生成的文本偏離預訓練模型的程度,防止出現 reward hacking 的問題。

我們最終的獎勵為以上三個獎勵的權重和,其中 KL 散度獎勵的系數是随着 KL 散度的大小動态變化的。想在 OpenRL 中使用該獎勵模型,使用者無需修改訓練代碼,隻需要在 nlp_ppo.yaml 檔案中添加 reward_class 參數即可:

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

OpenRL 支援使用者使用自定義的獎勵模型。首先,使用者需要編寫自定義獎勵模型 (需要繼承 BaseReward 類)。接着,使用者需要注冊自定義的獎勵模型,即在 train_ppo.py 添加以下代碼:

# train_ppo.py
from openrl.rewards.nlp_reward import CustomReward
from openrl.rewards import RewardFactory
RewardFactory.register ("CustomReward", CustomReward)           

最後,使用者隻需要在配置檔案中填寫自定義的獎勵模型即可:

reward_class:
    id: "CustomReward" # 自定義獎勵模型名稱
    args: {} # 使用者自定義獎勵函數可能用到的參數           

自定義訓練過程資訊輸出

OpenRL 還支援使用者自定義 wandb 和 tensorboard 的輸出内容。例如,在該任務的訓練過程中,我們還需要輸出各種類型獎勵的資訊和 KL 散度系數的資訊, 使用者可以在 nlp_ppo.yaml 檔案中加入 vec_info_class 參數來實作:

# nlp_ppo.yaml
vec_info_class:
    id: "NLPVecInfo" # 調用 NLPVecInfo 類以列印 NLP 任務中獎勵函數的資訊
# 設定 wandb 資訊
wandb_entity: openrl # 這裡用于指定 wandb 團隊名稱,請把 openrl 替換為你自己的團隊名稱
experiment_name: train_nlp # 這裡用于指定實驗名稱
run_dir: ./run_results/ # 這裡用于指定實驗資料儲存的路徑
log_interval: 1 # 這裡用于指定每隔多少個 episode 上傳一次 wandb 資料
# 自行填寫其他參數...           

修改完配置檔案後,在 train_ppo.py 檔案中啟用 wandb:

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

然後執行 python train_ppo.py –config nlp_ppo.yaml,稍後,便可以在 wandb 中看到如下的輸出:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

從上圖可以看到,wandb 輸出了各種類型獎勵的資訊和 KL 散度系數的資訊。

如果使用者還需要輸出其他資訊,還可以參考 NLPVecInfo 類 和 VecInfo 類來實作自己的 CustomVecInfo 類。然後,需要在 train_ppo.py 中注冊自定義的 CustomVecInfo 類:

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

最後,隻需要在 nlp_ppo.yaml 中填寫 CustomVecInfo 類即可啟用:

# nlp_ppo.yaml
vec_info_class:
    id: "CustomVecInfo" # 調用自定義 CustomVecInfo 類以輸出自定義資訊           

使用混合精度訓練加速

OpenRL 還提供了一鍵開啟混合精度訓練的功能。使用者隻需要在配置檔案中加入以下參數即可:

# nlp_ppo.yaml
use_amp: true # 開啟混合精度訓練           

對比評測

下表格展示了使用 OpenRL 訓練該對話任務的結果。結果顯示使用強化學習訓練後,模型各項名額皆有所提升。另外,從下表可以看出,相較于 RL4LMs , OpenRL 的訓練速度更快(在同樣 3090 顯示卡的機器上,速度提升 17% ),最終的性能名額也更好:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

最後,對于訓練好的智能體,使用者可以友善地通過 agent.chat () 接口進行對話:

# chat.py
from openrl.runners.common import ChatAgent as Agent
def 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 ()           

執行 python chat.py ,便可以和訓練好的智能體進行對話了:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

總結

OpenRL 架構經過了 OpenRL-Lab 的多次疊代并應用于學術研究和 AI 競賽,目前已經成為了一個較為成熟的強化學習架構。OpenRL-Lab 團隊将持續維護和更新 OpenRL,歡迎大家加入我們的開源社群,一起為強化學習的發展做出貢獻。更多關于 OpenRL 的資訊,可以參考:

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

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

緻謝

OpenRL 架構的開發吸取了其他強化學習架構的優點:

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 還處于持續開發和建設階段,未來 OpenRL 将會開源更多功能:

  • 支援智能體自博弈訓練
  • 加入離線強化學習、模範學習、逆強化學習算法
  • 加入更多強化學習環境和算法
  • 內建 Deepspeed 等加速架構
  • 支援多機分布式訓練

OpenRL Lab 團隊

OpenRL架構是由OpenRL Lab團隊開發,該團隊是第四範式公司旗下的強化學習研究團隊。第四範式長期緻力于強化學習的研發和工業應用。為了促進強化學習的産學研一體化,第四範式成立了OpenRL Lab研究團隊,目标是先進技術開源和人工智能前沿探索。成立不到一年,OpenRL Lab團隊已經在AAMAS發表過三篇論文,參加谷歌足球遊戲 11 vs 11比賽并獲得第三的成績。團隊提出的TiZero智能體,實作了首個從零開始,通過課程學習、分布式強化學習、自博弈等技術完成谷歌足球全場遊戲智能體的訓練:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

截止 2022 年 10 月 28 日,Tizero 在及第評測平台上排名第一:

訓練提速17%,第四範式強化學習研究架構,支援單、多智能體訓練

繼續閱讀