📘 JAX · 强化学习
30章 落地实战
⚡ 自动微分 · JIT · 向量化 · 函数式RL
🧩 友好色系
01
JAX与强化学习初探
定位
为什么是JAX？自动微分、JIT编译、向量化映射及其在RL中的定位
→
02
JAX基础入门
张量
张量操作、设备管理(CPU/GPU/TPU)、随机数生成与Key管理机制
→
03
JIT编译实战
@jit
@jit装饰器、静态/动态参数、Trace与Compile、常见编译陷阱
→
04
自动微分与vmap/pmap
变换
grad、vmap、pmap、lax.scan函数式变换深度解析与RL应用
→
05
RL环境向量化
Batch
使用JAX实现Batch环境 (Brax/EnvPool原理)，并行采样与数据收集
→
06
策略网络与价值网络
Flax/Haiku
纯函数式神经网络，参数管理、初始化与序列化
→
07
经验回放缓冲区
Replay
基于JAX固定大小循环缓冲区，高效采样与优先级采样
→
08
DQN的JAX实现
经典
从零搭建DQN，目标网络、ε-贪婪、损失函数与梯度更新
→
09
Double DQN & Dueling DQN
改进
改进算法在JAX中的实现，网络结构修改与训练稳定性提升
→
10
策略梯度方法
REINFORCE
REINFORCE算法实现，蒙特卡洛回报计算，基线技巧
→
11
Actor-Critic框架
A2C
A2C算法实现，优势函数估计，并行环境同步训练
→
12
PPO算法（上）
Clip/Penalty
PPO-Clip与PPO-Penalty原理，重要性采样与裁剪目标函数
→
13
PPO算法（下）
JAX实现
GAE优势估计，Mini-batch更新与Epoch训练
→
14
SAC算法
最大熵
最大熵强化学习，Soft Q函数与策略更新，自动温度调节
→
15
TD3算法
双延迟
双延迟DDPG，目标策略平滑正则化，JAX实现细节
→
16
分布式RL框架
pmap
数据并行训练，多设备梯度同步与参数广播
→
17
模型预测控制 MPC
可微分
基于JAX快速MPC，可微分动力学模型与规划
→
18
World Models
Dreamer
使用JAX训练世界模型，Dreamer架构初探与实现
→
19
MuZero算法
简化
基于JAX的MuZero简化实现，表示/动态/预测函数
→
20
多智能体强化学习
MADDPG/QMIX
JAX在多智能体中的应用，并行智能体训练
→
21
离线强化学习
CQL/BCQ
CQL、BCQ算法实现，JAX在离线数据上的高效处理
→
22
模仿学习与逆强化学习
GAIL/AIRL
生成对抗式模仿，JAX实现GAIL、AIRL算法
→
23
元强化学习
MAML-RL
MAML-RL、PEARL算法，JAX的grad变换在元学习中的应用
→
24
探索与利用
ICM/RND
ICM、RND、Curiosity-driven探索，JAX实现内在奖励计算
→
25
RL中的Transformer
决策Transformer
决策Transformer (DT) 架构，序列建模与JAX实现
→
26
可微分模拟器
物理引擎
使用JAX构建可微分物理引擎，基于梯度的策略优化
→
27
RL训练监控与调试
日志/保存
JAX下日志记录、模型保存/加载、TensorBoard集成
→
28
性能优化与部署
XLA/Serving
XLA编译优化、混合精度训练、模型导出与Serving
→
29
JAX RL生态工具
Brax/Rlax
Brax、Rlax、Chex、Optax等库的实战应用与最佳实践
→
30
综合项目实战
完整框架
从零构建完整JAX RL训练框架，涵盖环境、算法、训练与评估
→