TinyZero: 轻松训练AlphaZero风格的强化学习智能体

TinyZero是一个开源框架,旨在让研究人员和开发者能够轻松地在各种环境中训练AlphaZero风格的强化学习智能体。无论你是想复现经典的棋类游戏智能体,还是在自定义环境中探索新的应用,TinyZero都能为你提供强大而灵活的工具。

主要特性

• 易于使用:提供简洁的API,快速上手训练和评估智能体
• 灵活性强:支持自定义环境、模型和智能体
• 高效实现:优化的MCTS算法,支持GPU加速
• 可扩展性:方便添加新的游戏和任务
• 可视化:集成Weights & Biases支持,实时监控训练过程

快速入门

1. 安装依赖:

pip install -r requirements.txt

2. 训练智能体:

python3 tictactoe/two_dim/train.py

3. 评估智能体:

python3 tictactoe/two_dim/eval.py

你可以通过修改训练脚本中的参数来调整训练过程,如回合数、模拟次数等。

添加新环境

TinyZero的设计允许你轻松添加新的环境。你需要实现以下方法:

• reset(): 重置环境到初始状态
• step(action): 执行动作并更新环境状态
• get_legal_actions(): 返回合法动作列表
• undo_last_action(): 撤销上一个动作
• to_observation(): 将当前状态转换为观察值(numpy数组)
• get_result(): 返回游戏结果
• get_first_person_result(): 从当前玩家角度返回结果
• swap_result(result): 交换游戏结果

添加新模型

你可以通过实现以下方法来添加新的模型:

• __call__: 接收观察值,返回价值和策略
• value_forward(observation): 接收观察值,返回价值
• policy_forward(observation): 接收观察值,返回动作分布(策略)

添加新智能体

TinyZero支持添加各种类型的智能体。新智能体需要实现:

• value_fn(game): 接收游戏状态,返回价值估计
• policy_fn(game): 接收游戏状态,返回策略分布

在Google Colab中使用

TinyZero也可以在Google Colab等云平台上使用:

1. 安装wandb:

   !pip install wandb
   

2. 克隆仓库:

   !git clone https://github.com/s-casci/tinyzero.git
   

3. 训练(选择GPU运行时以加速):

   !cd tinyzero; python3 tictactoe/two_dim/train.py
   

4. 评估:

   !cd tinyzero; python3 tictactoe/two_dim/eval.py
   

TinyZero为强化学习研究和应用提供了一个灵活而强大的框架。无论你是想复现经典算法,还是探索新的应用领域,TinyZero都能成为你的得力助手。通过其清晰的接口和丰富的文档,你可以专注于算法和应用本身,而不必过多关注底层实现细节。

欢迎访问GitHub仓库了解更多信息,参与项目开发,或分享你基于TinyZero的研究成果!