Serendipity/rl-atari
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 改进DQN训练稳定性和性能

Browse Source 
- 将奖励裁剪替换为奖励缩放，保留奖励大小信号
- 添加学习率调度器，支持warmup和步进衰减
- 增加经验回放缓冲区大小至200,000
- 默认启用Dueling DQN架构
- 优化代码格式和参数传递
- 添加更多训练中间模型保存点
This commit is contained in:
刘航宇
2026-05-02 02:02:17 +08:00
parent 1c1cccd3f6
commit faf0d5ea42
12 changed files with 122 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
强化学习个人项目报告（Atari 游戏方向）/train.py
+58 -50
View File
@@ -1,6 +1,8 @@
"""Main training script for DQN on Space Invaders."""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import argparse
@@ -18,58 +20,61 @@ def main():
parser = argparse.ArgumentParser(description="DQN for Space Invaders")
# 环境参数
parser.add_argument("--env", type=str, default="ALE/SpaceInvaders-v5",
help="Atari环境ID")
parser.add_argument(
"--env", type=str, default="ALE/SpaceInvaders-v5", help="Atari环境ID"
)
# 训练参数
parser.add_argument("--steps", type=int, default=2_000_000,
help="总训练步数")
parser.add_argument("--lr", type=float, default=1e-4,
help="学习率")
parser.add_argument("--gamma", type=float, default=0.99,
help="折扣因子")
parser.add_argument("--batch-size", type=int, default=32,
help="批次大小")
parser.add_argument("--buffer-size", type=int, default=100_000,
help="经验回放缓冲区大小")
parser.add_argument("--steps", type=int, default=2_000_000, help="总训练步数")
parser.add_argument("--lr", type=float, default=1e-4, help="学习率")
parser.add_argument("--gamma", type=float, default=0.99, help="折扣因子")
parser.add_argument("--batch-size", type=int, default=32, help="批次大小")
parser.add_argument(
"--buffer-size", type=int, default=200_000, help="经验回放缓冲区大小"
)
# ε-greedy参数
parser.add_argument("--epsilon-start", type=float, default=1.0,
help="ε初始值")
parser.add_argument("--epsilon-end", type=float, default=0.01,
help="ε最终值")
parser.add_argument("--epsilon-decay", type=int, default=1_000_000,
help="ε衰减步数")
parser.add_argument("--epsilon-start", type=float, default=1.0, help="ε初始值")
parser.add_argument("--epsilon-end", type=float, default=0.01, help="ε最终值")
parser.add_argument(
"--epsilon-decay", type=int, default=1_000_000, help="ε衰减步数"
)
# 网络参数
parser.add_argument("--target-update", type=int, default=1000,
help="目标网络更新频率")
parser.add_argument("--double-dqn", action="store_true", default=True,
help="使用Double DQN")
parser.add_argument("--dueling", action="store_true", default=False,
help="使用Dueling DQN架构")
parser.add_argument(
"--target-update", type=int, default=500, help="目标网络更新频率"
)
parser.add_argument(
"--double-dqn", action="store_true", default=True, help="使用Double DQN"
)
parser.add_argument(
"--dueling", action="store_true", default=True, help="使用Dueling DQN架构"
)
# 学习率参数
parser.add_argument(
"--lr-decay-steps", type=int, default=1_000_000, help="学习率衰减步数"
)
parser.add_argument(
"--lr-decay-factor", type=float, default=0.5, help="学习率衰减因子"
)
parser.add_argument("--warmup-steps", type=int, default=10_000, help="预热步数")
# 评估参数
parser.add_argument("--eval-freq", type=int, default=10000,
help="评估频率")
parser.add_argument("--eval-episodes", type=int, default=10,
help="评估episode数")
parser.add_argument("--save-freq", type=int, default=50000,
help="模型保存频率")
parser.add_argument("--warmup", type=int, default=10000,
help="预热步数")
parser.add_argument("--eval-freq", type=int, default=10000, help="评估频率")
parser.add_argument("--eval-episodes", type=int, default=10, help="评估episode数")
parser.add_argument("--save-freq", type=int, default=50000, help="模型保存频率")
parser.add_argument("--warmup", type=int, default=10000, help="预热步数")
# 优先经验回放
parser.add_argument("--prioritized", action="store_true", default=False,
help="使用优先经验回放")
parser.add_argument(
"--prioritized", action="store_true", default=False, help="使用优先经验回放"
)
# 其他
parser.add_argument("--seed", type=int, default=42,
help="随机种子")
parser.add_argument("--save-dir", type=str, default="models",
help="模型保存目录")
parser.add_argument("--log-dir", type=str, default="logs",
help="日志目录")
parser.add_argument("--seed", type=int, default=42, help="随机种子")
parser.add_argument("--save-dir", type=str, default="models", help="模型保存目录")
parser.add_argument("--log-dir", type=str, default="logs", help="日志目录")
args = parser.parse_args()
@@ -110,14 +115,10 @@ def main():
# 创建经验回放缓冲区
if args.prioritized:
print("使用优先经验回放")
replay_buffer = PrioritizedReplayBuffer(
args.buffer_size, state_shape, device
)
replay_buffer = PrioritizedReplayBuffer(args.buffer_size, state_shape, device)
else:
print("使用标准经验回放")
replay_buffer = ReplayBuffer(
args.buffer_size, state_shape, device
)
replay_buffer = ReplayBuffer(args.buffer_size, state_shape, device)
# 创建智能体
agent = DQNAgent(
@@ -134,6 +135,9 @@ def main():
target_update_freq=args.target_update,
batch_size=args.batch_size,
double_dqn=args.double_dqn,
lr_decay_steps=args.lr_decay_steps,
lr_decay_factor=args.lr_decay_factor,
warmup_steps=args.warmup_steps,
)
# 创建训练器
@@ -146,20 +150,24 @@ def main():
eval_freq=args.eval_freq,
save_freq=args.save_freq,
num_eval_episodes=args.eval_episodes,
warmup_steps=args.warmup,
warmup_steps=args.warmup_steps,
)
# 打印配置
print("\n训练配置:")
print(f" 总步数: {args.steps:,}")
print(f" 学习率: {args.lr}")
print(f" 学习率衰减: 每{args.lr_decay_steps:,}步衰减{args.lr_decay_factor}倍")
print(f" Warmup步数: {args.warmup_steps:,}")
print(f" 折扣因子: {args.gamma}")
print(f" 批次大小: {args.batch_size}")
print(f" 缓冲区大小: {args.buffer_size:,}")
print(f" ε衰减: {args.epsilon_start} -> {args.epsilon_end} ({args.epsilon_decay:,}步)")
print(
f" ε衰减: {args.epsilon_start} -> {args.epsilon_end} ({args.epsilon_decay:,}步)"
)
print(f" 目标网络更新: 每{args.target_update}步")
print(f" Double DQN: {args.double_dqn}")
print(f" 预热步数: {args.warmup:,}")
print(f" Dueling DQN: {args.dueling}")
print("=" * 60)
# 开始训练
@@ -167,4 +175,4 @@ def main():
if __name__ == "__main__":
main()
main()