307 lines
11 KiB
Python
307 lines
11 KiB
Python
"""
|
||
DiscoRL ↔ Gym/SB3 Integration Guide
|
||
|
||
本文档总结了如何在 Stable-Baselines3 (SB3) 环境上使用 DiscoRL 的完整指南。
|
||
这是在自定义环境上部署的模板。
|
||
|
||
================================================================================
|
||
快速开始
|
||
================================================================================
|
||
|
||
1. 测试基础设施 (验证所有组件工作)
|
||
$ python scripts/test_disco_setup.py
|
||
|
||
2. 在 CartPole 上训练
|
||
$ python scripts/train_disco_cartpole.py
|
||
|
||
3. 验证集成 (完整的端到端测试)
|
||
$ python scripts/poc_integration.py
|
||
|
||
================================================================================
|
||
架构概述
|
||
================================================================================
|
||
|
||
DiscoRL 是一个 JAX/Haiku 框架，用于学习元学习的优化器。
|
||
要在 Gym/SB3 环境上使用它，我们需要：
|
||
|
||
1. 环境适配器 (disco_cartpole_env.py)
|
||
- 将 Gym 环境转换为 DiscoRL 的 Environment 接口
|
||
- 处理观测/动作的打包/解包
|
||
- 管理批量环境执行
|
||
|
||
2. 权重加载 (disco_weights.py)
|
||
- 加载预训练的 Disco103 元学习器权重
|
||
- 用于初始化元网络参数
|
||
|
||
3. 训练循环 (train_disco_cartpole.py)
|
||
- 数据收集 (rollout_trajectory)
|
||
- 参数更新 (agent.learner_step)
|
||
- 保存/加载检查点
|
||
|
||
================================================================================
|
||
核心组件说明
|
||
================================================================================
|
||
|
||
## 1. DiscoCartPoleEnv - 环境适配器
|
||
|
||
位置: scripts/disco_cartpole_env.py
|
||
|
||
关键方法:
|
||
- reset(rng_key=None) → (state, types.EnvironmentTimestep)
|
||
* 重置所有批量环境
|
||
* 返回初始观测作为 EnvironmentTimestep
|
||
|
||
- step(state, actions) → (state, types.EnvironmentTimestep)
|
||
* 执行动作，返回奖励/完成状态
|
||
* 自动处理已完成环境的重置
|
||
* 返回批量 EnvironmentTimestep
|
||
|
||
观测规格:
|
||
- Shape: (batch_size, 4) [CartPole 观测维度]
|
||
- Dtype: float32
|
||
|
||
动作规格:
|
||
- Type: Discrete(2) [Left=0, Right=1]
|
||
- Range: [0, 1]
|
||
|
||
## 2. DiscoRL Agent - 学习代理
|
||
|
||
主要操作:
|
||
|
||
a) 初始化
|
||
agent_settings = disco_agent.get_settings_disco()
|
||
agent = disco_agent.Agent(
|
||
single_observation_spec=obs_spec,
|
||
single_action_spec=act_spec,
|
||
agent_settings=agent_settings,
|
||
batch_axis_name=None,
|
||
)
|
||
|
||
b) 收集数据
|
||
actor_timestep, actor_state = agent.actor_step(
|
||
params=learner_state.params,
|
||
rng=rng,
|
||
timestep=env_timestep,
|
||
actor_state=actor_state,
|
||
)
|
||
→ 返回: 动作、策略输出等
|
||
|
||
c) 更新参数
|
||
new_learner_state, new_actor_state, logs = agent.learner_step(
|
||
rng=rng,
|
||
rollout=types.ActorRollout(...),
|
||
learner_state=learner_state,
|
||
agent_net_state=actor_state,
|
||
update_rule_params=meta_params,
|
||
is_meta_training=False, # 使用固定预训练的元网络
|
||
)
|
||
|
||
## 3. 数据流程
|
||
|
||
数据流动:
|
||
┌─────────────────────────┐
|
||
│ Gym 环境 (CartPole) │
|
||
└────────────┬────────────┘
|
||
│
|
||
↓
|
||
┌─────────────────────────────────────┐
|
||
│ DiscoCartPoleEnv.reset/step() │
|
||
│ (适配器，转换格式) │
|
||
└────────────┬────────────────────────┘
|
||
│
|
||
↓
|
||
┌─────────────────────────────────────┐
|
||
│ types.EnvironmentTimestep │
|
||
│ {observation, reward, step_type} │
|
||
└────────────┬────────────────────────┘
|
||
│
|
||
↓
|
||
┌─────────────────────────────────────┐
|
||
│ agent.actor_step() │
|
||
│ (策略推理) │
|
||
└────────────┬────────────────────────┘
|
||
│
|
||
↓
|
||
┌─────────────────────────────────────┐
|
||
│ types.ActorTimestep │
|
||
│ {actions, logits, agent_outs, ...} │
|
||
└────────────┬────────────────────────┘
|
||
│
|
||
┌─────────┴──────────┐
|
||
↓ ↓
|
||
[动作返回环境] [加入批量数据]
|
||
│ │
|
||
└─────────┬──────────┘
|
||
↓
|
||
┌─────────────────────────────────────┐
|
||
│ types.ActorRollout (堆叠轨迹) │
|
||
│ [T, B, ...] │
|
||
└────────────┬────────────────────────┘
|
||
│
|
||
↓
|
||
┌─────────────────────────────────────┐
|
||
│ agent.learner_step() │
|
||
│ (参数更新) │
|
||
└─────────────────────────────────────┘
|
||
|
||
其中:
|
||
T = 轨迹长度 (trajectory_length)
|
||
B = 批大小 (batch_size)
|
||
|
||
================================================================================
|
||
应用到自定义环境
|
||
================================================================================
|
||
|
||
要在自己的环境 (如 gym_env_250326_erase.py) 上使用 DiscoRL:
|
||
|
||
1. 创建环境适配器
|
||
创建文件: scripts/disco_custom_env.py
|
||
|
||
from disco_cartpole_env import DiscoCartPoleEnv
|
||
import your_env # 导入自定义环境
|
||
|
||
class DiscoCustomEnv(DiscoCartPoleEnv):
|
||
def __init__(self, batch_size: int = 1):
|
||
# 不要调用 super().__init__()
|
||
# 创建自定义环境实例而不是 CartPole
|
||
self._envs = [your_env.create_env() for _ in range(batch_size)]
|
||
|
||
# 根据自定义环境构建规格
|
||
base_env = self._envs[0]
|
||
obs_space = base_env.observation_space
|
||
act_space = base_env.action_space
|
||
|
||
# 创建 dm_env 规格
|
||
from dm_env import specs
|
||
self._single_observation_spec = {...} # 基于自定义环境
|
||
self._single_action_spec = {...} # 基于自定义环境
|
||
|
||
# 保持其余逻辑相同
|
||
|
||
2. 调整观测/动作处理
|
||
- 确保观测转换为 float32 JAX 数组
|
||
- 确保动作转换为正确的类型 (int 或 float，取决于动作空间)
|
||
|
||
3. 更新训练配置
|
||
在 train_disco_*.py 中:
|
||
env = DiscoCustomEnv(batch_size=4)
|
||
# 使用相同的训练循环
|
||
|
||
================================================================================
|
||
关键设计决策
|
||
================================================================================
|
||
|
||
1. CPU-Only JAX
|
||
• 原因: 避免 GPU 内存冲突
|
||
• 设置: os.environ['JAX_PLATFORMS'] = 'cpu' 在脚本顶部
|
||
|
||
2. 离散动作处理
|
||
• CartPole 已经有离散动作 (0/1)
|
||
• 不需要连续→离散映射
|
||
• 直接通过 int 值到 Gym
|
||
|
||
3. 批量执行
|
||
• 所有操作在 Python 级别批处理 (不使用 jax.vmap)
|
||
• 维持简单性和通用性
|
||
|
||
4. 元学习禁用
|
||
• is_meta_training=False
|
||
• 使用预初始化的元参数 (不学习优化器)
|
||
• 目标: 学习环境特定参数
|
||
|
||
================================================================================
|
||
文件结构
|
||
================================================================================
|
||
|
||
scripts/
|
||
├── disco_cartpole_env.py # CartPole ↔ DiscoRL 适配器 [核心]
|
||
├── disco_weights.py # 权重加载工具 [辅助]
|
||
├── train_disco_cartpole.py # 训练循环 [示例]
|
||
├── test_disco_setup.py # 完整测试 [验证]
|
||
├── poc_integration.py # 端到端 PoC [演示]
|
||
├── DISCO_RL_GUIDE.py # 详细文档 [参考]
|
||
└── [将来]
|
||
├── train_disco_custom.py # 适应自定义环境
|
||
└── disco_custom_env.py # 自定义环境适配器
|
||
|
||
================================================================================
|
||
已测试的组件
|
||
================================================================================
|
||
|
||
✓ DiscoCartPoleEnv
|
||
- 批量重置
|
||
- 批量步进
|
||
- 自动恢复已完成的环境
|
||
- 正确的 EnvironmentTimestep 格式
|
||
|
||
✓ DiscoRL Agent
|
||
- 初始化学习者/执行者状态
|
||
- actor_step() 推理
|
||
- learner_step() 参数更新
|
||
- 梯度计算和优化器步骤
|
||
|
||
✓ 完整训练循环
|
||
- 轨迹收集
|
||
- 批量数据聚合
|
||
- 学习器步骤
|
||
- 保存检查点
|
||
|
||
✓ 与 CartPole 兼容性
|
||
- 离散动作空间 (0/1)
|
||
- 连续观测 (4D)
|
||
- 标准奖励信号
|
||
|
||
================================================================================
|
||
故障排除
|
||
================================================================================
|
||
|
||
问题: "JAX GPU 内存错误"
|
||
解决: 在文件顶部添加 os.environ['JAX_PLATFORMS'] = 'cpu'
|
||
|
||
问题: "ActorRollout 字段错误"
|
||
解决: 不包括 'behaviour_agent_out'，只使用 'agent_outs'
|
||
|
||
问题: "CartPole 步进警告"
|
||
解决: 在 DiscoCartPoleEnv.step() 中使用 _episode_done 标志来防止
|
||
在完成后重新步进
|
||
|
||
问题: "权重加载失败"
|
||
解决: 目前省略预训练权重加载
|
||
只使用随机初始化的元参数
|
||
可以手动复制权重（超出范围）
|
||
|
||
================================================================================
|
||
下一步
|
||
================================================================================
|
||
|
||
1. 使用 DiscoRL 在自定义环境上训练
|
||
|
||
a) 复制 disco_cartpole_env.py → disco_custom_env.py
|
||
b) 调整环境创建逻辑
|
||
c) 复制 train_disco_cartpole.py → train_disco_custom.py
|
||
d) 更新环境导入
|
||
e) 运行: python scripts/train_disco_custom.py
|
||
|
||
2. 比较与 SB3 基线
|
||
|
||
查看 eval_disco_vs_sb3.py (框架已准备)
|
||
实现权重加载以进行真实的预训练评估
|
||
|
||
3. 调整超参数
|
||
|
||
batch_size: 环境并行化程度
|
||
trajectory_length: 学习器的展开长度
|
||
learning_rate: 优化器学习率
|
||
num_iterations: 训练步骤数
|
||
|
||
4. 监控训练
|
||
|
||
跟踪: average_reward, total_loss
|
||
绘制: 奖励曲线，损失曲线
|
||
比较: DiscoRL vs. 标准 RL
|
||
|
||
================================================================================
|
||
"""
|
||
|
||
print(__doc__)
|