第四范式開源強化學習研究通用框架，支持單智能體、多智能體訓練，還可訓練自然語言任務！訓練速度提升17%

2023-05-05 17:49:23

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

51CTO讀者成長計劃社群招募，咨詢小助手（微信號：CTOjishuzhan）

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

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

目前，OpenRL已經在GitHub開源：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訓練智能體只需要簡單的四步：創建環境 => 初始化模型 => 初始化智能體 => 開始訓練！

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

此外，對于多智能體、自然語言等任務的訓練，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,
        asynchrnotallow=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訓練前后智能體的表現：

二、加載配置文件

此外，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 來可視化訓練過程：

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)：

五、訓練自然語言對話任務

最近的研究表明，強化學習也可以用于訓練語言模型，并且能顯著提升模型的性能。目前，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,asynchrnotallow=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,asynchrnotallow=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中看到如下的輸出：

從上圖可以看到，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類以輸出自定義信息

10、使用混合精度訓練加速

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

對比評測

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

下表格展示了使用OpenRL訓練該對話任務的結果。結果顯示使用強化學習訓練后，模型各項指標皆有所提升。另外，從下表可以看出，相較于 RL4LMs ， OpenRL的訓練速度更快（在同樣3090顯卡的機器上，速度提升 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 ，便可以和訓練好的智能體進行對話了：

十一、總結

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等加速框架
支持多機分布式訓練

責任編輯：武曉燕來源： 51CTO技術棧

開源 OpenRL 自然語言