Fix/replay buffer nan loss

optimization_plan.md

@@ -0,0 +1,101 @@
+# Rainbow DQN 损失优化计划
+
+## 问题分析
+
+当前训练中观察到的问题：
+- **平均损失值过高**：2.6769（正常范围应为 0.1-0.5）
+- **Q值预测不稳定**：TD误差较大，表明网络学习效果不佳
+- **收敛速度慢**：需要更多回合才能达到理想性能
+
+## 优化策略
+
+### 阶段一：超参数调优（立即执行）
+
+#### 1.1 学习率优化
+- **当前值**：1e-4
+- **建议值**：5e-5 或 3e-5
+- **原因**：降低学习率可以减少Q值更新的波动，提高训练稳定性
+
+#### 1.2 批量大小调整
+- **当前值**：32
+- **建议值**：64 或 128
+- **原因**：更大的批量可以提供更稳定的梯度估计
+
+#### 1.3 目标网络更新频率
+- **当前值**：1000步
+- **建议值**：2000步
+- **原因**：减少目标网络更新频率，提高训练稳定性
+
+### 阶段二：损失函数改进
+
+#### 2.1 使用 Huber Loss
+- **当前**：平均绝对误差 (MAE)
+- **建议**：Huber Loss（对异常值更鲁棒）
+- **实现**：修改 `_compute_standard_loss` 函数
+
+#### 2.2 梯度裁剪优化
+- **当前值**：10
+- **建议值**：1.0 或 0.5
+- **原因**：更严格的梯度裁剪防止梯度爆炸
+
+### 阶段三：网络结构优化
+
+#### 3.1 权重初始化改进
+- 使用 Xavier 或 He 初始化
+- 确保初始Q值在合理范围内
+
+#### 3.2 噪声网络参数调整
+- **sigma_init**：从 0.4 降低到 0.2
+- 减少噪声强度，提高学习稳定性
+
+### 阶段四：训练策略优化
+
+#### 4.1 预热训练
+- 前1000步使用更小的学习率
+- 逐步增加到目标学习率
+
+#### 4.2 学习率调度
+- 实现指数衰减或余弦退火
+- 在训练后期进一步降低学习率
+
+## 实施计划
+
+### 第一步：创建优化配置文件
+创建 `optimized_config.py` 包含所有优化参数
+
+### 第二步：修改损失函数
+在 `agent.py` 中实现 Huber Loss
+
+### 第三步：调整训练脚本
+修改 `train.py` 支持新的优化参数
+
+### 第四步：测试验证
+运行短期训练验证优化效果
+
+## 预期效果
+
+- **损失值**：从 2.6+ 降低到 0.3-0.8
+- **收敛速度**：提升 30-50%
+- **最终性能**：奖励提升 15-25%
+- **训练稳定性**：显著改善
+
+## 监控指标
+
+1. **平均损失趋势**：应持续下降
+2. **Q值分布**：应趋于稳定
+3. **TD误差**：应逐渐减小
+4. **奖励方差**：应降低
+5. **梯度范数**：应保持稳定
+
+## 风险评估
+
+- **过度保守**：学习率过低可能导致收敛过慢
+- **批量过大**：可能导致内存不足
+- **参数冲突**：多个优化可能相互影响
+
+## 回滚策略
+
+如果优化效果不佳：
+1. 逐步回退参数到原始值
+2. 单独测试每个优化项
+3. 保留有效的优化，移除无效的

optimized_config.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""
+优化配置文件 - Rainbow DQN 损失优化参数
+
+本文件包含经过优化的超参数设置，旨在降低训练损失并提高模型性能。
+使用方法：python src/train.py --config optimized_config.py
+"""
+
+# 基础训练参数
+OPTIMIZED_CONFIG = {
+    # 核心超参数优化
+    'lr': 3e-5,                    # 降低学习率：从 1e-4 到 3e-5
+    'batch_size': 64,              # 增加批量大小：从 32 到 64
+    'target_update': 2000,         # 增加目标网络更新间隔：从 1000 到 2000
+    'gamma': 0.99,                 # 保持折扣因子不变
+
+    # 缓冲区和经验回放
+    'buffer_size': 100000,         # 保持缓冲区大小
+    'prioritized_replay': True,    # 启用优先经验回放
+
+    # 探索策略（非噪声网络）
+    'epsilon_start': 1.0,
+    'epsilon_final': 0.01,
+    'epsilon_decay': 150000,       # 增加衰减步数，更平缓的探索
+
+    # Rainbow DQN 特有参数
+    'use_noisy': True,             # 启用噪声网络
+    'use_distributional': True,    # 启用分布式Q学习
+    'n_atoms': 51,                 # 分布原子数
+    'v_min': -10,                  # 值函数范围
+    'v_max': 10,
+
+    # 自适应 N-step 参数优化
+    'base_n_step': 3,
+    'max_n_step': 8,               # 降低最大N步：从 10 到 8
+    'adapt_n_step_freq': 1500,     # 增加调整频率：从 1000 到 1500
+    'td_error_threshold_low': 0.08, # 降低阈值：从 0.1 到 0.08
+    'td_error_threshold_high': 0.4, # 降低阈值：从 0.5 到 0.4
+
+    # 状态增强参数
+    'use_state_augmentation': True,
+    'aug_noise_scale': 3.0,        # 降低噪声强度：从 5.0 到 3.0
+
+    # 训练控制
+    'episodes': 2000,              # 增加训练回合数
+    'save_interval': 50,           # 更频繁保存：从 100 到 50
+    'eval_interval': 10,           # 保持评估间隔
+
+    # 网络结构优化参数
+    'noisy_sigma_init': 0.2,       # 降低噪声初始化：从 0.4 到 0.2
+    'grad_clip_norm': 1.0,         # 更严格的梯度裁剪：从 10 到 1.0
+
+    # 学习率调度参数
+    'use_lr_scheduler': True,      # 启用学习率调度
+    'lr_decay_factor': 0.95,       # 学习率衰减因子
+    'lr_decay_steps': 10000,       # 每10000步衰减一次
+    'min_lr': 1e-6,                # 最小学习率
+
+    # 损失函数优化
+    'use_huber_loss': True,        # 启用 Huber Loss
+    'huber_delta': 1.0,            # Huber Loss 的 delta 参数
+
+    # 预热训练参数
+    'warmup_steps': 1000,          # 预热步数
+    'warmup_lr_factor': 0.1,       # 预热期学习率因子
+
+    # 监控和日志
+    'log_loss_freq': 100,          # 每100步记录详细损失
+    'log_grad_norm': True,         # 记录梯度范数
+    'log_q_values': True,          # 记录Q值统计
+}
+
+# 验证配置
+def validate_config(config):
+    """
+    验证配置参数的合理性
+    """
+    assert config['lr'] > 0, "学习率必须为正数"
+    assert config['batch_size'] > 0, "批量大小必须为正数"
+    assert 0 < config['gamma'] <= 1, "折扣因子必须在(0,1]范围内"
+    assert config['v_min'] < config['v_max'], "值函数范围设置错误"
+    assert config['base_n_step'] <= config['max_n_step'], "N步参数设置错误"
+    print("✓ 配置验证通过")
+
+# 配置比较函数
+def compare_with_default():
+    """
+    与默认配置进行比较
+    """
+    default_config = {
+        'lr': 1e-4,
+        'batch_size': 32,
+        'target_update': 1000,
+        'max_n_step': 10,
+        'grad_clip_norm': 10,
+        'noisy_sigma_init': 0.4,
+    }
+
+    print("配置对比：")
+    for key in default_config:
+        if key in OPTIMIZED_CONFIG:
+            print(f"{key}: {default_config[key]} → {OPTIMIZED_CONFIG[key]}")
+
+if __name__ == "__main__":
+    validate_config(OPTIMIZED_CONFIG)
+    compare_with_default()
+    print("\n优化配置已准备就绪！")

quick_test_optimization.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+快速优化测试脚本
+用于验证Rainbow DQN优化配置的效果
+"""
+
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+os.environ["QT_QPA_PLATFORM_PLUGIN_PATH"] = ""
+
+import torch
+import numpy as np
+import time
+import argparse
+from torch.utils.tensorboard import SummaryWriter
+
+from src.model import RainbowDQN
+from src.agent import RainbowAgent
+from src.utils import make_env
+from optimized_config import OPTIMIZED_CONFIG
+
+def quick_test():
+    """
+    快速测试优化配置效果
+    """
+    print("=== Rainbow DQN 优化配置快速测试 ===")
+
+    # 使用优化配置
+    config = OPTIMIZED_CONFIG.copy()
+    config['episodes'] = 3  # 只测试3个回合
+    config['memory_size'] = 1000  # 减小内存大小
+    config['learning_starts'] = 100  # 更早开始学习
+    config['use_per'] = config.get('prioritized_replay', True)
+    config['augmentation_config'] = {'add_noise': {'scale': config.get('aug_noise_scale', 3.0)}} if config.get('use_state_augmentation', True) else None
+
+    print(f"学习率: {config['lr']}")
+    print(f"批量大小: {config['batch_size']}")
+    print(f"梯度裁剪: {config['grad_clip_norm']}")
+    print(f"Huber Delta: {config['huber_delta']}")
+    print("="*50)
+
+    # 设备
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"使用设备: {device}")
+
+    # 创建环境
+    env = make_env("ALE/Assault-v5")
+
+    # 获取环境信息
+    state_shape = env.observation_space.shape
+    n_actions = env.action_space.n
+
+    print(f"状态形状: {state_shape}")
+    print(f"动作数量: {n_actions}")
+
+    # 创建模型
+    model = RainbowDQN(
+        input_shape=state_shape,
+        n_actions=n_actions,
+        use_noisy=config['use_noisy'],
+        use_distributional=config['use_distributional'],
+        n_atoms=config['n_atoms'],
+        v_min=config['v_min'],
+        v_max=config['v_max']
+    ).to(device)
+
+    # 创建目标模型
+    target_model = RainbowDQN(
+        input_shape=state_shape,
+        n_actions=n_actions,
+        use_noisy=config['use_noisy'],
+        use_distributional=config['use_distributional'],
+        n_atoms=config['n_atoms'],
+        v_min=config['v_min'],
+        v_max=config['v_max']
+    ).to(device)
+
+    # 创建智能体
+    agent = RainbowAgent(
+        model=model,
+        target_model=target_model,
+        env=env,
+        device=device,
+        base_n_step=config['base_n_step'],
+        max_n_step=config['max_n_step'],
+        adapt_n_step_freq=config['adapt_n_step_freq'],
+        td_error_threshold_low=config['td_error_threshold_low'],
+        td_error_threshold_high=config['td_error_threshold_high'],
+        augmentation_config=config['augmentation_config'],
+        use_noisy=config['use_noisy'],
+        use_distributional=config['use_distributional'],
+        n_atoms=config['n_atoms'],
+        v_min=config['v_min'],
+        v_max=config['v_max'],
+        huber_delta=config['huber_delta'],
+        grad_clip_norm=config['grad_clip_norm'],
+        buffer_size=config['memory_size'],
+        batch_size=config['batch_size'],
+        gamma=config['gamma'],
+        lr=config['lr'],
+        epsilon_start=config['epsilon_start'],
+        epsilon_final=config['epsilon_final'],
+        epsilon_decay=config['epsilon_decay'],
+        target_update=config['target_update'],
+        prioritized_replay=config['use_per']
+    )
+
+    print("开始快速测试...")
+
+    # 训练循环
+    episode_rewards = []
+    episode_losses = []
+
+    for episode in range(config['episodes']):
+        state, _ = env.reset()
+        episode_reward = 0
+        episode_loss = []
+        steps = 0
+        start_time = time.time()
+
+        while True:
+            # 选择动作
+            action = agent.select_action(state)
+
+            # 执行动作
+            next_state, reward, done, truncated, _ = env.step(action)
+
+            # 存储经验
+            agent.store_experience(state, action, reward, next_state, done or truncated)
+
+            # 学习
+            if len(agent.memory) >= config['learning_starts']:
+                loss = agent.update_model()
+                if loss is not None:
+                    episode_loss.append(loss)
+
+            episode_reward += reward
+            state = next_state
+            steps += 1
+
+            if done or truncated or steps >= 200:  # 限制最大步数
+                break
+
+        # 计算平均损失
+        avg_loss = np.mean(episode_loss) if episode_loss else 0.0
+        duration = time.time() - start_time
+
+        episode_rewards.append(episode_reward)
+        episode_losses.append(avg_loss)
+
+        print(f"QUICK TEST (Ep: {episode+1}/{config['episodes']}) "
+              f"Reward: {episode_reward:.2f}, Avg Loss: {avg_loss:.4f}, "
+              f"Duration: {duration:.2f}s, Steps: {steps}")
+
+    # 输出测试结果
+    print("\n=== 测试结果 ===")
+    print(f"平均奖励: {np.mean(episode_rewards):.2f}")
+    print(f"平均损失: {np.mean(episode_losses):.4f}")
+    print(f"损失标准差: {np.std(episode_losses):.4f}")
+
+    # 检查损失趋势
+    if len(episode_losses) >= 2:
+        loss_trend = episode_losses[-1] - episode_losses[0]
+        if loss_trend < 0:
+            print(f"✓ 损失下降趋势: {loss_trend:.4f}")
+        else:
+            print(f"⚠ 损失上升趋势: {loss_trend:.4f}")
+
+    print("\n=== 优化配置验证完成 ===")
+
+    env.close()
+    return episode_rewards, episode_losses
+
+if __name__ == "__main__":
+    quick_test()