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