DDPG InvertedPendulum-v5 Demo
使用纯 PyTorch 实现 DDPG,在 MuJoCo 的 InvertedPendulum-v5 环境中训练一个连续控制智能体。
项目简介
DDPG(Deep Deterministic Policy Gradient)是一种经典的 Actor-Critic 深度强化学习算法,专为连续动作空间设计。这个项目默认展示的是 InvertedPendulum-v5,适合作为学习 MuJoCo 连续控制强化学习的起点。
虽然脚本仍然支持通过 --env-id 切换到其他 MuJoCo 环境,但这篇项目页的默认案例就是 InvertedPendulum-v5。相比已经单独拆分出来并做过针对性调参的 ddpg-reacher-v5、ddpg-pusher-v5,这里更像是一个从简单任务切入的 DDPG baseline 示例。
算法原理
DDPG 结合了 DPG(确定性策略梯度)和 DQN 的核心思想:
┌─────────────────────────────────────────────┐
│ DDPG 架构 │
│ │
│ Actor(策略网络)──→ 输出确定性动作 a = μ(s) │
│ Critic(价值网络)──→ 评估 Q(s, a) │
│ Target 网络 ──→ 稳定训练目标 │
│ Replay Buffer ──→ 打破样本相关性 │
└─────────────────────────────────────────────┘
核心更新公式
Critic 更新(最小化 TD 误差):
Actor 更新(最大化 Q 值):
Target 网络软更新:
代码结构
ddpg-invertedpendulum-v5/
├── figs/
│ └── invertedpendulum-success.gif
├── index.mdx
└── train_ddpg_invertedpendulum_v5.py
本项目刻意将所有代码放在单个文件中,方便阅读和理解。训练脚本当前约 333 行,主要模块包括:
| 模块 | 说明 |
|---|---|
ReplayBuffer | 高效的 NumPy 环形经验回放缓冲区 |
MLP | 通用多层感知机基础网络 |
Actor | 策略网络,输出经 tanh 缩放到动作范围的确定性动作 |
Critic | Q 值网络,输入 (state, action) 拼接向量 |
DDPGConfig | 使用 dataclass 管理所有超参数 |
DDPGAgent | 封装训练逻辑:动作选择、网络更新、软更新 |
train() | 主训练循环,含评估与模型保存 |
代码详解
经验回放缓冲区
使用预分配的 NumPy 数组实现环形缓冲区,避免 Python list 的内存碎片和 GC 开销:
class ReplayBuffer:
def __init__(self, obs_dim: int, action_dim: int, capacity: int) -> None:
self.capacity = capacity
self.obs = np.zeros((capacity, obs_dim), dtype=np.float32)
self.actions = np.zeros((capacity, action_dim), dtype=np.float32)
self.rewards = np.zeros((capacity, 1), dtype=np.float32)
self.next_obs = np.zeros((capacity, obs_dim), dtype=np.float32)
self.dones = np.zeros((capacity, 1), dtype=np.float32)
self.ptr = 0
self.size = 0
Actor 网络
Actor 输出经过 tanh 激活后,通过线性缩放映射到环境的实际动作范围 [low, high]:
class Actor(nn.Module):
def forward(self, obs: torch.Tensor) -> torch.Tensor:
raw_action = torch.tanh(self.backbone(obs))
return raw_action * self.action_scale + self.action_bias
探索策略
训练初期使用随机采样(start_steps),之后在确定性动作上叠加高斯噪声进行探索:
if step <= config.start_steps:
action = env.action_space.sample() # 纯随机探索
else:
action = agent.act(obs, noise_scale=0.1) # 高斯噪声探索
Truncation 处理
对于因时间限制截断的 episode,仍然需要 bootstrap 目标值(即 done=False),只有真正终止的状态才设为 done=True:
replay_done = float(terminated) # 只用 terminated,忽略 truncated
环境要求
pip install gymnasium[mujoco] torch numpy tqdm imageio
如果你本地已经有仓库里的 conda 环境,推荐直接使用:
source /home/jj/anaconda3/etc/profile.d/conda.sh
conda activate easy-robotics
:::tip 硬件说明
MuJoCo 环境本身不需要 GPU,但 PyTorch 训练可以利用 CUDA 加速。脚本默认 --device auto 会自动检测。
:::
快速开始
默认命令会在 InvertedPendulum-v5 上训练,并把 actor 权重保存到 artifacts/ddpg_invertedpendulum_v5_actor.pt:
python projects/ddpg-invertedpendulum-v5/train_ddpg_invertedpendulum_v5.py
如果你想基于这份实现切换到别的 MuJoCo 环境:
python projects/ddpg-invertedpendulum-v5/train_ddpg_invertedpendulum_v5.py \
--env-id HalfCheetah-v5 \
--total-steps 100000 \
--start-steps 5000 \
--hidden-dim 256 \
--batch-size 256 \
--gamma 0.99 \
--tau 0.005 \
--exploration-noise 0.1 \
--save-path artifacts/ddpg_halfcheetah_v5_actor.pt
如果你想按“两步走”的方式先保存 checkpoint,再加载模型录制 gif:
python projects/ddpg-invertedpendulum-v5/train_ddpg_invertedpendulum_v5.py \
--save-path artifacts/ddpg_invertedpendulum_v5_actor.pt
python projects/ddpg-invertedpendulum-v5/train_ddpg_invertedpendulum_v5.py \
--skip-train \
--load-path artifacts/ddpg_invertedpendulum_v5_actor.pt \
--record-path projects/ddpg-invertedpendulum-v5/figs/invertedpendulum-success.gif
如果你想看已经拆分好的调参版项目:
Reacher-v5:/projects/ddpg-reacher-v5Pusher-v5:/projects/ddpg-pusher-v5
默认参数
| 参数 | 默认值 | 说明 |
|---|---|---|
--env-id | InvertedPendulum-v5 | Gymnasium 环境 ID |
--total-steps | 20000 | 总训练步数 |
--start-steps | 1000 | 纯随机探索步数 |
--batch-size | 256 | 每次更新的 batch 大小 |
--buffer-size | 200000 | 经验回放缓冲区容量 |
--gamma | 0.99 | 折扣因子 |
--tau | 0.005 | 目标网络软更新系数 |
--actor-lr | 1e-3 | Actor 学习率 |
--critic-lr | 1e-3 | Critic 学习率 |
--hidden-dim | 256 | 隐藏层维度 |
--exploration-noise | 0.1 | 探索噪声比例 |
--seed | 42 | 随机种子 |
--eval-interval | 2000 | 每隔多少步评估一次 |
--eval-episodes | 5 | 每次评估的 episode 数 |
--load-path | None | 加载已经训练好的 actor checkpoint |
--skip-train | False | 只加载模型并执行录制,不重新训练 |
--record-path | None | 导出 gif 或 mp4 |
--record-episodes | 1 | 录制多少个 episode |
--record-fps | 20 | 导出媒体帧率 |
--record-frame-skip | 8 | 每隔多少环境步采一帧 |
--save-path | artifacts/ddpg_invertedpendulum_v5_actor.pt | 模型保存路径 |
--device | auto | 设备选择:auto / cpu / cuda |
实现要点
- Actor 网络输出经过
tanh映射回环境动作范围 - Critic 网络输入为
(obs, action)拼接后的向量 - 训练初期使用纯随机探索,之后在确定性动作上叠加高斯噪声
- 时间截断
truncated不会被当作真正终止,目标值仍然继续 bootstrap - 可以先训练生成 actor checkpoint,再通过
--load-path+--skip-train单独录制 gif - 这个项目默认聚焦
InvertedPendulum-v5,也可以继续作为其他 MuJoCo 任务的起点脚本
最终效果
下面这段 gif 使用默认 seed=42 训练保存 checkpoint 后,再单独加载模型录制得到:

训练输出示例
device=cuda env=InvertedPendulum-v5
training: 10%|██ | 2000/20000 [00:12] episode_return=87.0 avg10=62.3 actor_loss=0.012 critic_loss=0.035
[eval] step=2000 avg_return=78.40
training: 20%|████ | 4000/20000 [00:24] episode_return=256.0 avg10=189.5 actor_loss=0.008 critic_loss=0.021
[eval] step=4000 avg_return=312.60
...
saved_actor=artifacts/ddpg_invertedpendulum_v5_actor.pt
final_eval_return=1000.00
适合进一步尝试的方向
- TD3:在 DDPG 基础上加入 clipped double-Q、延迟更新和目标策略平滑,显著提升稳定性
- SAC:引入最大熵框架,自动调节探索程度,通常是连续控制的默认选择
- 多环境对比:尝试
Hopper-v5、Walker2d-v5、Ant-v5等更复杂的环境 - 超参数搜索:对
hidden_dim、学习率、tau等进行网格搜索或贝叶斯优化