""" 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__)