训练提速17%，第四范式开源强化学习研究框架，支持单、多智能体训练

pip install openrl

conda install -c openrl 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
env = make ("CartPole-v1", env_num=9) # 创建环境，并设置环境并行数为 9
net = Net (env) # 创建神经网络
agent = Agent (net) # 初始化智能体
agent.train (total_time_steps=20000) # 开始训练，并设置环境运行总步数为 20000

# 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 ()

# 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

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

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

# 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 ()

# 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 ()

# 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

# 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
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 ()

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

# nlp_ppo.yaml
reward_class:
    id: NLPReward # 奖励模型名称
    args: {
        # 用于意图判断的模型的名称或路径
        "intent_model": rajkumarrrk/roberta-daily-dialog-intent-classifier,
        # 用于计算 KL 散度的预训练模型的名称或路径
        "ref_model": roberta-base, # 用于意图判断的 tokenizer 的名称或路径
    }

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

reward_class:
    id: "CustomReward" # 自定义奖励模型名称
    args: {} # 用户自定义奖励函数可能用到的参数

# 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
agent.train (total_time_steps=100000, use_wandb=True)

# train_ppo.py # 注册自定义输出信息类 
VecInfoFactory.register ("CustomVecInfo", CustomVecInfo)

# nlp_ppo.yaml
vec_info_class:
    id: "CustomVecInfo" # 调用自定义 CustomVecInfo 类以输出自定义信息

# nlp_ppo.yaml
use_amp: true # 开启混合精度训练

# 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 ()