前言
从零开始学习ai文章系列计划是个人在《动手学深度学习》和《磨菇书》两本书的学习中的个人笔记,文章也会以课本中的章节分开,即每个章节一片笔记。我会尽量的把主要内容以及遇到的难点进行记录与解决,如果哪里有错误的欢迎指正。或者不清晰的可以直接查看原文部分。
《蘑菇书》原文(课本):https://datawhalechina.github.io/easy-rl/#/
(由于有时候公式太多,可能会直接贴图片)
蘑菇书的文章结构不会跟之前《动手学深度学习》按照原文章节进行,个人会适当调节。
(第十章实验鸽了,好奇心实验做的道心破碎,本章是蘑菇书学习最后一章)
0. 实验环境
实验在 MiniGrid-DoorKey-8x8-v0 或
MiniGrid-DoorKey-5x5-v0 环境中进行。
下图是比较大的一个MiniGrid环境,我们以此为例说明环境的参数
1 env = gym.make("MiniGrid-DoorKey-16x16-v0" , render_mode="human" )
观测空间
1 2 3 4 5 6 7 8 9 { 'image' : array(7x7x3), 'direction' : 0 , 'mission' : 'use the key to open the door and then get to the goal' }
image → 视野中的地图
direction → 朝向
mission → 任务描述
Agent 只能看到前方 7×7 的区域。
形状:7 × 7 × 3
每个格子包含 3个值:(object_type, color, state)
MiniGrid 渲染时:
empty cell 不画任何物体
背景颜色就是 黑色
所以[1,0,0] 这里看到是黑色的背景。
很多人会误以为 floor 是可走地面,但实际上:
empty:普通可走格子
floor:特殊地面对象(有些环境才用)
动作空间
终止条件:
捡起钥匙,打开门后,进入到绿色方块处。
超过最大步数。max_steps = 10 * size * size,我们环境中size =
16,所以:max_steps = 2560
收集专家经验
人工玩MiniGrid-DoorKey游戏,记录每一步的(s,a),当作专家数据。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 DATA_DIR = Path("data" ) DATA_DIR.mkdir(exist_ok=True ) ENV_ID = "MiniGrid-DoorKey-5x5-v0" TRAJ_FILE = DATA_DIR / "trajectories.pkl" ACTION_LEFT = 0 ACTION_RIGHT = 1 ACTION_FORWARD = 2 ACTION_PICKUP = 3 ACTION_TOGGLE = 5 ACTION_NAMES = {0 :"left" ,1 :"right" ,2 :"forward" ,3 :"pickup" ,4 :"drop" ,5 :"toggle" ,6 :"done" } N_ACTIONS = 7 def collect_demos (): """pygame 渲染,人工键盘操作,记录 (state, action, next_state, done) 轨迹""" env = gym.make(ENV_ID, render_mode="human" ) obs, _ = env.reset() trajectories = [] if TRAJ_FILE.exists(): with open (TRAJ_FILE, "rb" ) as f: trajectories = pickle.load(f) print (f"载入已有轨迹 {len (trajectories)} 条" ) current_traj = [] episode = 0 total_saved = len (trajectories) print ("\n=== 人工演示收集 ===" ) print ("← → 转向 | ↑ 前进 | Space 拾取/开门 | R 重置 | S 保存 | Q 退出" ) clock = pygame.time.Clock() running = True while running: env.render() action = None for event in pygame.event.get(): if event.type == pygame.QUIT: running = False elif event.type == pygame.KEYDOWN: if event.key == pygame.K_LEFT: action = ACTION_LEFT elif event.key == pygame.K_RIGHT: action = ACTION_RIGHT elif event.key == pygame.K_UP: action = ACTION_FORWARD elif event.key == pygame.K_SPACE: _inner = env while hasattr (_inner, "env" ): _inner = _inner.env fwd = _inner.front_pos cell = _inner.grid.get(*fwd) if cell and cell.type == "door" : action = ACTION_TOGGLE else : action = ACTION_PICKUP elif event.key == pygame.K_r: discarded = len (current_traj) obs, _ = env.reset() current_traj = [] print (f" 重置 (丢弃 {discarded} 步) | 已存 {total_saved} 条轨迹" ) elif event.key == pygame.K_s: if current_traj: success = current_traj[-1 ][4 ] trajectories.append({ "episode" : total_saved, "length" : len (current_traj), "success" : success, "transitions" : current_traj }) total_saved += 1 with open (TRAJ_FILE, "wb" ) as f: pickle.dump(trajectories, f) print (f" [已保存] 轨迹 #{total_saved} : {len (current_traj)} 步, " f"{'成功 ✓' if success else '未完成' } → {TRAJ_FILE} " ) obs, _ = env.reset() current_traj = [] episode += 1 else : print (" 当前无轨迹数据,请先完成操作再按 S" ) elif event.key == pygame.K_q: running = False if action is not None : state = extract_state(obs, env) next_obs, reward, terminated, truncated, _ = env.step(action) done = terminated or truncated next_state = extract_state(next_obs, env) current_traj.append((state, action, next_state, float (reward), done)) obs = next_obs if done: success = reward > 0 print (f" 局结束: {len (current_traj)} 步, reward={reward:.2 f} " f" {'✓ 成功! 按 S 保存,或按 R 丢弃重来' if success else '✗ 超时,按 R 重置' } " ) clock.tick(10 ) env.close() print (f"\n共保存 {total_saved} 条轨迹 → {TRAJ_FILE} " ) return trajectories
我们要讲一下extract_state函数,因为后续我们也会用到。
我们从obs中,结合世界环境提取了以下几点基本信息
1 2 3 4 5 6 agent_dir, carrying, key_relative, door_relative, door_open, dist_to_key,
实际上我们不应该结合世界坐标,因为这是未知的,但为了训练的方便我们也放进去了。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 def extract_state (obs, env ): """ 将 MiniGrid 观测 + 环境内部状态 → 固定维度特征向量 (float32) 维度: 7×7×3 展平(图像) + 标量特征 = 147 + 8 = 155 标量特征(8维): [0] agent_dir : 朝向 (0~3 归一化到 0~1) [1] carrying : 是否持有钥匙 (0/1) [2-3] key_relative : 钥匙相对 agent 的 (dx, dy),除以 width 归一化 持有钥匙时置 (0, 0) [4-5] door_relative : 门相对 agent 的 (dx, dy),除以 width 归一化 [6] door_open : 门是否已开 (0/1) [7] dist_to_key : 到钥匙的曼哈顿距离,除以 width 归一化 持有钥匙时置 0(辅助信号,帮助学习"靠近") 为什么去掉 agent 绝对坐标: 绝对坐标在 5x5 和 8x8 中分布不同,导致 g_θ 网络在迁移时 收到训练分布外的输入,输出混乱的奖励信号。 相对坐标在任意地图尺寸下语义一致。 """ image = obs["image" ].astype(np.float32) / 10.0 image_flat = image.flatten() inner = env while hasattr (inner, "env" ): inner = inner.env width = float (inner.width) agent_pos = np.array(inner.agent_pos, dtype=np.float32) agent_dir = np.array([inner.agent_dir / 3.0 ], dtype=np.float32) carrying = np.array([1.0 if inner.carrying is not None else 0.0 ], dtype=np.float32) door_pos_abs = None key_pos_abs = None door_open = np.array([0.0 ], dtype=np.float32) for x in range (inner.width): for y in range (inner.height): cell = inner.grid.get(x, y) if cell is None : continue if cell.type == "door" : door_pos_abs = np.array([x, y], dtype=np.float32) door_open = np.array([1.0 if cell.is_open else 0.0 ], dtype=np.float32) elif cell.type == "key" : key_pos_abs = np.array([x, y], dtype=np.float32) if key_pos_abs is not None : key_relative = (key_pos_abs - agent_pos) / width else : key_relative = np.array([0.0 , 0.0 ], dtype=np.float32) if door_pos_abs is not None : door_relative = (door_pos_abs - agent_pos) / width else : door_relative = np.array([0.0 , 0.0 ], dtype=np.float32) if key_pos_abs is not None : dist_to_key = np.array( [np.abs (key_relative).sum ()], dtype=np.float32 ) else : dist_to_key = np.array([0.0 ], dtype=np.float32) scalars = np.concatenate([ agent_dir, carrying, key_relative, door_relative, door_open, dist_to_key, ]) return np.concatenate([image_flat, scalars]).astype(np.float32)
1. 行为克隆 (Behavior Cloning)
专家在方块环境中,演示如何移动到钥匙、捡起钥匙、打开门、到达目标。将这些步骤中的
(s,a) 收集起来,当作我们的专家数据。
我们把专家的动作a当作标签,让 agent 在相同状态s
时,也输出正确的动作a 。
与分类问题的计算一样
1 2 logits = policy(states) loss = F.cross_entropy(logits, actions)
专家的动作a当作标签,为目标类别。我们让该类别可能性越来越大。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 def train_bc (epochs=100 , batch_size=256 , lr=1e-3 ): if not TRAJ_FILE.exists(): print ("无轨迹文件, 请先 collect" ); return with open (TRAJ_FILE, "rb" ) as f: trajs = pickle.load(f) states, actions = [], [] for traj in trajs: for (s, a, ns, r, d) in traj["transitions" ]: states.append(s) actions.append(a) states = torch.FloatTensor(np.array(states)) actions = torch.LongTensor(actions) N = len (states) print (f"BC 数据: {N} 条 (s,a) 对, 来自 {len (trajs)} 条轨迹" ) policy = PolicyNet() opt = Adam(policy.parameters(), lr=lr, weight_decay=1e-4 ) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs) loss_history = [] best_loss = float ("inf" ) for epoch in range (1 , epochs + 1 ): idx = torch.randperm(N) total_loss = 0.0 for i in range (0 , N, batch_size): b = idx[i:i+batch_size] logits = policy(states[b]) loss = F.cross_entropy(logits, actions[b]) opt.zero_grad(); loss.backward(); opt.step() total_loss += loss.item() * len (b) total_loss /= N sched.step() loss_history.append(total_loss) if total_loss < best_loss: best_loss = total_loss torch.save(policy.state_dict(), BC_MODEL) if epoch % 10 == 0 : acc = (policy(states).argmax(-1 ) == actions).float ().mean().item() print (f" Epoch {epoch:4d} /{epochs} | loss={total_loss:.4 f} | " f"train_acc={acc:.3 f} | lr={sched.get_last_lr()[0 ]:.5 f} " ) _plot_loss(loss_history, "BC Training Loss" , "bc_loss.png" ) print (f"\nBC 模型已保存 → {BC_MODEL} " ) return policy, loss_history
行为克隆只是单纯的动作模仿,不会有任何逻辑。
2. 逆强化学习(inverse
reinforcement leanring)
MaxEnt
IRL(最大熵逆强化学习) 是一种从专家演示中推断奖励函数的算法 。
牢记:从专家演示 中推断奖励函数 的算法。因为后面理论公式一堆,很容易不知道自己干啥去了。
2.1 理论
我们先详细说明奖励网络的理论。
首先,MaxEnt IRL 的目标是最大化专家轨迹的对数似然
(即找一个θ参数的模型,传入专家轨迹链τ,最大化返回概率值):
D是我们的专家经验。 \[
Target(\theta) = \sum_{\tau \in D}\log P(\tau|\theta)
\] 其中专家轨迹τ的概率计算如下 \[
P(\tau|\theta) = \exp(\theta^{\top}f(\tau))/Z(\theta)
\] 对于公式中的θ和f,理解如下
理论上奖励是 \(R(τ)=θ^⊤f(τ)\) ,需要手动设计特征
f ,再学习权重
θ 。但我们代码直接用神经网络替代了这整个过程。
好,回到刚才我们的概率计算中,我们最终目的是对最大化专家轨迹的对数似然进行求导
\[
P(\tau|\theta) = \exp(\theta^{\top}f(\tau))/Z(\theta)
\] 其中Z为所有可能轨迹的概率求和,使我们的 概率P值归纳到0-1之间。
\[
Z(θ) = \sum_τ exp(θ^{\top} f(τ))
\] 现在我们有 N 条专家轨迹,目标是让这些轨迹在参数 θ
下出现的概率尽量大。 \[
Target(\theta) = \sum_{\tau \in D}\log P(\tau|\theta)=\sum_{\tau \in
D}log \frac{exp(θ^{\top} f(τ))}{Z(θ)}
\] 首先拆开log,得到 \[
=\sum_{\tau \in D}[θ^{\top} f(τ)-logZ(\theta)]
\] 然后提取常数项,logZ(θ) 与具体哪条轨迹 τ 无关,对 N
条轨迹求和就出现 N 次: \[
=θ^{\top}\sum_{\tau \in D}f(\tau)-N*logZ(\theta)
\] 除以N(不影响我们求最优解),定义 \[
\tilde{f}=\frac{1}{N}\sum_{\tau \in D}f(\tau)
\] 我门的目标变成如下 \[
Traget(\theta) = θ^{\top} \tilde{f} - log Z(\theta)
\] 最后我们对参数θ开始求导。
第一项 \(θ^{\top} \tilde{f}\) 对 θ
求导,得到 \(\tilde{f}\)
第二项 \(\log Z(\theta)\)
用链式法则: \[
\nabla_\theta\log Z(\theta) = \frac{1}{Z(\theta)}\nabla_\theta Z(\theta)
\] 我们记得 \[
Z(θ) = \sum_τ exp(θ^{\top} f(τ))
\] 因此 \[
\nabla_\theta Z(θ) = \sum_τ exp(θ^{\top} f(τ))*f(τ)
\] 最终我们得到 \[
\nabla_\theta\log Z(\theta) = \sum_τ\frac{ exp(θ^{\top}
f(τ))}{Z(\theta)}*f(τ)
\] 我们知道前面P的定义如下 \[
P(\tau|\theta) = \exp(\theta^{\top}f(\tau))/Z(\theta)
\] 因此 \[
\nabla_\theta\log Z(\theta) = \sum_τP(\tau|\theta)
*f(τ)=\mathbb{E}_{P(\tau|\theta)}[f(\tau)]
\] 最终两项合并,得到 \[
\nabla_\theta
Target(\theta)=\tilde{f}-\mathbb{E}_{P(\tau|\theta)}[f(\tau)]
\]
其中
\(\tilde{f}\)
是专家的轨迹链的特征均值,来自数据,因此固定不变。
\(\mathbb{E}_{P(\tau|\theta)}[f(\tau)]\)
表示:在当前奖励参数 \(\theta\)
下,由策略分布 \(P(\tau|\theta)\)
产生的轨迹,其特征的期望。即固定奖励网络的参数θ,用最优策略(走奖励最大的一条路)走过的所有轨迹,对他们的特征值
f(τ) 求期望。
固定当前奖励函数后,智能体会倾向于选择 reward 更高的轨迹,因此 \(\mathbb{E}_{P(\tau|\theta)}[f(\tau)]\)
可以理解为:
“当前奖励函数诱导出的策略,会产生什么样的轨迹特征”。
理论上是这样,但实际上算不出来,"解出最优策略跑无数条轨迹"在连续空间 里根本算不了——轨迹是无限的。
所以我们实现的代码中用了一个近似:
用当前智能体实际跑出来的轨迹,来近似"当前奖励下最优策略的轨迹"。
我们用当前 agent
输出的动作概率,跑出m条轨迹链,然后放到奖励函数中求均值,就是我们的
\(\mathbb{E}_{P(\tau|\theta)}[f(\tau)]\)
近似。
在经典 MaxEnt IRL 中: \[
R(τ)=w^⊤f(τ)
\] 其中:
\(f(\tau)\) 是人为设计的 feature
vector;
\(w\) 是 feature 的线性权重。
例如: \[
f(\tau)= \begin{bmatrix} \text{靠近目标}\\ \text{撞墙次数}\\ \text{步数}
\end{bmatrix}
\] reward 本质上是:对 feature vector
的线性投影(projection) 。即: \(\mathbb R^d \rightarrow \mathbb R\) 。
因此 \(f(\tau)\) 是多维向量,而
\(R(\tau) = w^\top f(\tau)\) 则是
feature 在方向 \(w\)
上的投影值,也就是最终 reward。
因此,我们也可以将原始的 feature matching: \[
\tilde{f}-\mathbb{E}_{P(\tau|\theta)}[f(\tau)]
\] 理解为:“比较 expert 与 agent 在 feature space
中的差异 ”。
如果对两边同时左乘 \(w^\top\)
,则有: \[
w^\top \left( \tilde f - \mathbb E_\pi[f(\tau)] \right)
\] 根据线性性质: \[
= w^\top \tilde f - w^\top \mathbb E_\pi[f(\tau)]
\] 进一步得到: \[
= \mathbb E_{expert}[R(\tau)] - \mathbb E_\pi[R(\tau)]
\] 现在我们的 θ 指的是 reward network 的参数。
此时:
都已经被神经网络隐式学习,因此不再显式构造 \(f(\tau)\)
更准确地说,现在实际上变成了: \[
R_\theta(s) = w^\top \phi_\theta(s)
\] 其中:
\(\phi_\theta(s)\)
是网络内部学习到的 hidden feature;
\(w\) 是最后一层线性映射。
因此定义专家状态上的平均 reward: \[
R_e = R(\tau).mean() \tag{τ是所有的专家轨迹链}
\] 因为奖励r 仅与状态s有关, reward_net
是逐状态打分的,因此我们也能够写成 \[
R_e = \sum_tR(s_t).mean()
\]
然后固定当前 agent 参数,用 agent 采样得到 m
条轨迹,收集其中状态。对应的平均 reward 为: \[
R_a = R(\tau).mean()=\sum_tR(s_t).mean() \tag{τ是agent收集的轨迹链}
\] 于是我们最终得到 reward network 的优化目标: \[
∇Target≈R_e−R_a=R(e_s).mean()−R(a_s).mean()
\] 由于我们的目标是最大化该目标,而 PyTorch
默认执行最小化,因此定义 loss: \[
loss=−(R(e_s).mean()−R(a_s).mean())
\]
这个推导里面最有意思的一点是概率与奖励的关系,即 \[
P(\tau|\theta) = \exp(\theta^{\top}f(\tau))/Z(\theta)
\]
正是有该转换的基础,我们才能将奖励与概率之间的相互影响写出来。因此如果我们想探索
结果A与影响因子B之间的关系,我们可以去找找现有的公式或者自行构造一个合理的公式。
2.2 代码实验
首先从专家经验中,把状态都拿出来,转化成tensor
1 2 3 4 5 6 expert_states = [] for traj in trajs: for (s, a, ns, r, d) in traj["transitions" ]: expert_states.append(s) expert_states = torch.FloatTensor(np.array(expert_states))
定义我们需要的网络
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 class RewardNet (nn.Module): """IRL 奖励网络: state → scalar reward""" def __init__ (self, state_dim=STATE_DIM, hidden=128 ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1 ) ) def forward (self, x ): return self .net(x).squeeze(-1 ) class PolicyNet (nn.Module): """策略网络: state → action logits""" def __init__ (self, state_dim=STATE_DIM, hidden=256 , n_actions=N_ACTIONS ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, n_actions) ) def forward (self, x ): return self .net(x) def act (self, state_np, greedy=False ): with torch.no_grad(): s = torch.FloatTensor(state_np).unsqueeze(0 ) logits = self .forward(s) if greedy: return logits.argmax(dim=-1 ).item() dist = Categorical(logits=logits) return dist.sample().item() class ValueNet (nn.Module): """价值网络: state → scalar (用于 PPO/IRL)""" def __init__ (self, state_dim=STATE_DIM, hidden=256 ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1 ) ) def forward (self, x ): return self .net(x).squeeze(-1 ) reward_net = RewardNet() policy = PolicyNet() value_net = ValueNet()
其中,我们的重点是奖励网络reward_net。policy是我们的agent,value_net则是进行更新的时候计算gae用的。
接下来我们开始每轮的训练,首先agent收集数据
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 def _collect_rollout (env, policy, reward_net, n_steps ): """用 IRL reward 收集轨迹""" obs, _ = env.reset() rollout = [] for _ in range (n_steps): s = extract_state(obs, env) a = policy.act(s) next_obs, _, term, trunc, _ = env.step(a) done = term or trunc with torch.no_grad(): r = reward_net(torch.FloatTensor(s).unsqueeze(0 )).item() rollout.append((s, a, extract_state(next_obs, env), r, done)) obs = next_obs if done: obs, _ = env.reset() return rollout for irl_ep in range (1 , irl_epochs + 1 ): rollout = _collect_rollout(env, policy, reward_net, ppo_steps) agent_states = torch.FloatTensor(np.array([t[0 ] for t in rollout]))
能看到 a =
policy.act(s),policy也是用当前的概率分布获取动作a,并非直接获取最大概率的动作a。
也许你也注意到,我们数据的收集并不是每走过一个episode存一次,而是直接
for _ in range(n_steps):
,这样我们每轮就固定收集到了n_steps步数的数据,如果提前结束的画,就会重置继续执行。
if done: obs, _ = env.reset()
备齐了材料后,我们开始奖励网络的训练。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 for irl_ep in range (1 , irl_epochs + 1 ): rollout = _collect_rollout(env, policy, reward_net, ppo_steps) agent_states = torch.FloatTensor(np.array([t[0 ] for t in rollout])) if best_sr < 0.9 : r_opt.zero_grad() r_expert = reward_net(expert_states).mean() r_agent = reward_net(agent_states).mean() irl_loss = -(r_expert - r_agent) irl_loss.backward() torch.nn.utils.clip_grad_norm_(reward_net.parameters(), 0.5 ) r_opt.step() irl_loss_hist.append(irl_loss.item()) else : irl_loss_hist.append(irl_loss_hist[-1 ] if irl_loss_hist else 0.0 )
这里和我们之间的loss一样 \[
loss = -(R(states\_e).mean() - R(states\_a).mean())
\] 最后就是更新我们当前的policy
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 def _ppo_update (policy, value, opt, rollout, epochs, gamma=0.99 , clip=0.2 , lam=0.95 ): """PPO clip 更新""" states = torch.FloatTensor(np.array([t[0 ] for t in rollout])) actions = torch.LongTensor([t[1 ] for t in rollout]) rewards = [t[3 ] for t in rollout] dones = [t[4 ] for t in rollout] values = value(states).detach().numpy() advs, returns = [], [] gae = 0.0 for i in reversed (range (len (rewards))): nv = values[i + 1 ] if i + 1 < len (values) else 0.0 mask = 0.0 if dones[i] else 1.0 delta = rewards[i] + gamma * nv * mask - values[i] gae = delta + gamma * lam * mask * gae advs.insert(0 , gae) returns.insert(0 , gae + values[i]) advs = torch.FloatTensor(advs) returns = torch.FloatTensor(returns) advs = (advs - advs.mean()) / (advs.std() + 1e-8 ) with torch.no_grad(): old_logp = Categorical(logits=policy(states)).log_prob(actions) N = len (states) for _ in range (epochs): idx = torch.randperm(N) for i in range (0 , N, 256 ): b = idx[i:i+256 ] logits = policy(states[b]) dist = Categorical(logits=logits) logp = dist.log_prob(actions[b]) ratio = (logp - old_logp[b]).exp() surr1 = ratio * advs[b] surr2 = ratio.clamp(1 -clip, 1 +clip) * advs[b] v_pred = value(states[b]) loss = (-torch.min (surr1, surr2).mean() + 0.5 * F.mse_loss(v_pred, returns[b]) - 0.01 * dist.entropy().mean()) opt.zero_grad(); loss.backward() torch.nn.utils.clip_grad_norm_( list (policy.parameters()) + list (value.parameters()), 0.5 ) opt.step() return loss.item() for irl_ep in range (1 , irl_epochs + 1 ): rollout = _collect_rollout(env, policy, reward_net, ppo_steps) agent_states = torch.FloatTensor(np.array([t[0 ] for t in rollout])) if best_sr < 0.9 : r_opt.zero_grad() r_expert = reward_net(expert_states).mean() r_agent = reward_net(agent_states).mean() irl_loss = -(r_expert - r_agent) irl_loss.backward() torch.nn.utils.clip_grad_norm_(reward_net.parameters(), 0.5 ) r_opt.step() irl_loss_hist.append(irl_loss.item()) else : irl_loss_hist.append(irl_loss_hist[-1 ] if irl_loss_hist else 0.0 ) _ppo_update(policy, value_net, p_opt, rollout, ppo_epochs) avg_r = np.mean([t[3 ] for t in rollout]) reward_hist.append(avg_r)
虽然内容有点多,我们慢慢来。首先要介绍GAE
GAE:在advantage熵做TD(λ)
资格迹:在 value/Q 上做TD(λ)
这是 PPO 中的 GAE-λ公式为:
代码解释如下
如果不太清楚代码中的gae计算,可以参考如下举例
advs 是 advantage(优势函数),数学的定义是:
那么如果我们想计算Q值,直接通过 advs+v[i] 算回来即可。
即代码中的:returns.insert(0, gae + values[i])
接下来计算初始策略的log值,为了后续PPO中的计算权重使用
1 2 3 4 with torch.no_grad(): old_logp = Categorical(logits=policy(states)).log_prob(actions)
接下来就是PPO的正常更新了
可以回顾下我们PPO的目标导数公式
3.
对抗学习(Generative Adversarial Imitation Learning)
3.1 理论介绍
GAIL(Generative Adversarial Imitation Learning)= 对抗式模仿学习
核心思想借鉴自 Generative Adversarial Network:
有两个网络在“对抗”:
策略 π(Generator):产生行为(像专家一样行动)
判别器 D(Discriminator):判断行为是“专家的”还是“模型的”
判别器任务是学会区分专家和智能体,可以输入专家的(s,a)与agent的(s,a)做交叉熵
1 2 loss_d = (F.binary_cross_entropy(d_expert, torch.full_like(d_expert, 0.9 )) + F.binary_cross_entropy(d_agent, torch.zeros_like(d_agent)))
在收集数据过程中,我们将判别器的结果修改成奖励 \[
r_{GAIL}=\log D(s,a)−\log (1−D(s,a))
\]
当前动作越接近专家,则奖励越高。即,gail_r
把"像专家的行为"映射成高奖励。
Policy 通过 PPO
更新,趋向高奖励路径。最终判别器无法区分两者,是收敛的标志。
收敛时 D(s,a)≈0.5 。此时 \[
r_{GAIL}=\log0.5−\log0.5=0
\]
收敛意味着智能体行为 已经和专家一样,判别器"放弃区分",此时奖励信号自然归零——任务完成了,不再需要额外激励 。
3.2 代码
首先还是从专家经验中,提取我们需要的内容 (s, a)
1 2 3 4 5 6 7 8 expert_sa_states = [] expert_sa_actions = [] for traj in trajs: for (s, a, ns, r, d) in traj["transitions" ]: expert_sa_states.append(s) expert_sa_actions.append(a) E_s = torch.FloatTensor(np.array(expert_sa_states)) E_a = F.one_hot(torch.LongTensor(expert_sa_actions), N_ACTIONS).float ()
创建我们需要的网络。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 class DiscriminatorNet (nn.Module): """GAIL 判别器: (state, action) → P(expert)""" def __init__ (self, state_dim=STATE_DIM, n_actions=N_ACTIONS, hidden=128 ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim + n_actions, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1 ), nn.Sigmoid() ) def forward (self, states, actions_onehot ): x = torch.cat([states, actions_onehot], dim=-1 ) return self .net(x).squeeze(-1 ) class PolicyNet (nn.Module): """策略网络: state → action logits""" def __init__ (self, state_dim=STATE_DIM, hidden=256 , n_actions=N_ACTIONS ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, n_actions) ) def forward (self, x ): return self .net(x) def act (self, state_np, greedy=False ): with torch.no_grad(): s = torch.FloatTensor(state_np).unsqueeze(0 ) logits = self .forward(s) if greedy: return logits.argmax(dim=-1 ).item() dist = Categorical(logits=logits) return dist.sample().item() class ValueNet (nn.Module): """价值网络: state → scalar (用于 PPO/IRL)""" def __init__ (self, state_dim=STATE_DIM, hidden=256 ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1 ) ) def forward (self, x ): return self .net(x).squeeze(-1 ) disc = DiscriminatorNet() policy = PolicyNet() value = ValueNet()
Policy和Value网络和之前一样。我们的判别器D需要输出是否是专家的概率,因此最后一层输出1,且需要经过sigmoid函数激活映射到
(0,1) 之间。
接下来开始每轮的更新,首先还是用当前agent收集轨迹链,将收集到的 (s,a)
转换成tensor。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 def _collect_rollout_gail (env, policy, disc, n_steps ): """用 GAIL reward = log D(s,a) - log(1-D(s,a)) 收集轨迹""" obs, _ = env.reset() rollout = [] for _ in range (n_steps): s = extract_state(obs, env) a = policy.act(s) next_obs, _, term, trunc, _ = env.step(a) done = term or trunc with torch.no_grad(): s_t = torch.FloatTensor(s).unsqueeze(0 ) a_oh = F.one_hot(torch.tensor([a]), N_ACTIONS).float () d = disc(s_t, a_oh).item() d = np.clip(d, 1e-6 , 1 - 1e-6 ) r = np.log(d) - np.log(1 - d) rollout.append((s, a, extract_state(next_obs, env), r, done)) obs = next_obs if done: obs, _ = env.reset() return rollout for ep in range (1 , gail_epochs + 1 ): rollout = _collect_rollout_gail(env, policy, disc, ppo_steps) A_s = torch.FloatTensor(np.array([t[0 ] for t in rollout])) A_a = F.one_hot(torch.LongTensor([t[1 ] for t in rollout]), N_ACTIONS).float ()
有了专家和agent数据,我们可以开始更新判别器了。逻辑很简单,用交叉熵做损失函数,专家数据经过判别器的结果标签为1,agent数据经过判别器的结果标签为0。计算损失然后进行更新。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 for ep in range (1 , gail_epochs + 1 ): rollout = _collect_rollout_gail(env, policy, disc, ppo_steps) A_s = torch.FloatTensor(np.array([t[0 ] for t in rollout])) A_a = F.one_hot(torch.LongTensor([t[1 ] for t in rollout]), N_ACTIONS).float () n_disc_updates = 3 for _ in range (n_disc_updates): idx_e = torch.randperm(len (E_s))[:min (512 , len (E_s))] idx_a = torch.randperm(len (A_s))[:min (512 , len (A_s))] d_expert = disc(E_s[idx_e], E_a[idx_e]) d_agent = disc(A_s[idx_a], A_a[idx_a]) loss_d = (F.binary_cross_entropy(d_expert, torch.ones_like(d_expert)) + F.binary_cross_entropy(d_agent, torch.zeros_like(d_agent))) d_opt.zero_grad(); loss_d.backward(); d_opt.step() disc_loss_hist.append(loss_d.item())
接下来是对agent的更新。和之前一样,也是先计算优势和回报然后进行PPO更新。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 def _ppo_update (policy, value, opt, rollout, epochs, gamma=0.99 , clip=0.2 , lam=0.95 ): """PPO clip 更新""" states = torch.FloatTensor(np.array([t[0 ] for t in rollout])) actions = torch.LongTensor([t[1 ] for t in rollout]) rewards = [t[3 ] for t in rollout] dones = [t[4 ] for t in rollout] values = value(states).detach().numpy() advs, returns = [], [] gae = 0.0 for i in reversed (range (len (rewards))): nv = values[i + 1 ] if i + 1 < len (values) else 0.0 mask = 0.0 if dones[i] else 1.0 delta = rewards[i] + gamma * nv * mask - values[i] gae = delta + gamma * lam * mask * gae advs.insert(0 , gae) returns.insert(0 , gae + values[i]) advs = torch.FloatTensor(advs) returns = torch.FloatTensor(returns) advs = (advs - advs.mean()) / (advs.std() + 1e-8 ) with torch.no_grad(): old_logp = Categorical(logits=policy(states)).log_prob(actions) N = len (states) for _ in range (epochs): idx = torch.randperm(N) for i in range (0 , N, 256 ): b = idx[i:i+256 ] logits = policy(states[b]) dist = Categorical(logits=logits) logp = dist.log_prob(actions[b]) ratio = (logp - old_logp[b]).exp() surr1 = ratio * advs[b] surr2 = ratio.clamp(1 -clip, 1 +clip) * advs[b] v_pred = value(states[b]) loss = (-torch.min (surr1, surr2).mean() + 0.5 * F.mse_loss(v_pred, returns[b]) - 0.01 * dist.entropy().mean()) opt.zero_grad(); loss.backward() torch.nn.utils.clip_grad_norm_( list (policy.parameters()) + list (value.parameters()), 0.5 ) opt.step() return loss.item() for ep in range (1 , gail_epochs + 1 ): rollout = _collect_rollout_gail(env, policy, disc, ppo_steps) A_s = torch.FloatTensor(np.array([t[0 ] for t in rollout])) A_a = F.one_hot(torch.LongTensor([t[1 ] for t in rollout]), N_ACTIONS).float () n_disc_updates = 3 for _ in range (n_disc_updates): idx_e = torch.randperm(len (E_s))[:min (512 , len (E_s))] idx_a = torch.randperm(len (A_s))[:min (512 , len (A_s))] d_expert = disc(E_s[idx_e], E_a[idx_e]) d_agent = disc(A_s[idx_a], A_a[idx_a]) loss_d = (F.binary_cross_entropy(d_expert, torch.ones_like(d_expert)) + F.binary_cross_entropy(d_agent, torch.zeros_like(d_agent))) d_opt.zero_grad(); loss_d.backward(); d_opt.step() disc_loss_hist.append(loss_d.item()) _ppo_update(policy, value, p_opt, rollout, ppo_epochs) if ep % 10 == 0 : sr = _eval_policy(env, policy, 20 ) sr_hist.append(sr) print (f" GAIL ep {ep:3d} /{gail_epochs} | " f"disc_loss={loss_d.item():.4 f} | success_rate={sr:.2 %} " )
4.
对抗逆强化学习(Adversarial Inverse Reinforcement Learning)
4.1 理论
在之前的girl中,我们的判别器仅仅用来判断是否是专家数据。比如输入专家数据
(s,a) ,判别器就将输出往1方向拟合。如果是agent数据
(s,a),则往0方向拟合。
GAIL 的判别器就是一个纯粹的二分类器: \[
D(s,a)=σ(f(s,a))
\] f(s,a)
是任意神经网络,没有任何结构约束,只负责分类 ,学到的东西无法解释 。
在 AIRL 中,AIRL
给判别器加了结构约束 ,强制它的形式为: \[
D_\theta(s,a,s')=\frac{\exp(f_\theta(s,a,s'))}{\exp(f_\theta(s,a,s'))+\pi(a|s)}
\] 其中\(f_\theta\)
被进一步分解为:
我们将 f 强制约束成 优势计算 的结构。
我们判别器D的训练目标和 GAIL 一样,做二分类:
对于专家数据的(s,a,s'),将其拟合成1,而对于agent轨迹的(s,a,s'),将其拟合为0.
1 2 3 4 loss_d = ( F.binary_cross_entropy(d_expert, torch.ones_like(d_expert)) + F.binary_cross_entropy(d_agent, torch.zeros_like(d_agent)) )
当收敛时,策略 π 完全模仿专家时,判别器无法区分两者,此时: \[
D_\theta(s,a,s')=\frac{\exp(f_\theta(s,a,s'))}{\exp(f_\theta(s,a,s'))+\pi(a|s)}=0.5
\] 即判断为agent和专家的概率都为50%,无法区分开来。
此时 \[
\exp(f)=π(a|s)
\] 即 \[
f(s,a,s')=\logπ(a|s)
\] 为什么这意味着 \(f→A^∗\)
最优策略满足: \[
π^∗(a∣s)∝exp(A^∗(s,a))
\] 即,在最优策略中,策略在某状态s下的所有动作a概率,与
对应优势的exp值成正比 \(exp(A^∗(s,a))\)
实际上:
举例:
接下来对 \(\pi^*(a|s)\) 两边取
log:
即: \[
\logπ^∗(a∣s)=A^∗(s,a)+C(s)
\] 其中 C(s) 是只和状态有关的归一化常数,会被 \(h_\phi(s)\) 吸收。
所以收敛时: \[
f^∗(s,a,s')=\logπ^∗(a|s)=A^∗(s,a)
\] 因为我们强制 \(f=g_θ(s,a)+γh_ϕ(s')−h_ϕ(s)\) ,
因此一组可行解是: \[
g_θ=r^∗,h_ϕ=V^∗
\] 但该解不是唯一的,因为 \(h_\phi(s)\)
只刻画“状态之间的相对差异”,而不是绝对数值。只要同时对 reward 和 h
做一致的调整(相当于给状态加一个势能函数),最终结果不会发生变化。
总之,最终我们学到了我们需要的奖励函数 \(g_\theta\)
,它真正解释了我们做哪些动作有什么奖励。且与环境无关。(奖励函数的值并一定和真实奖励相同,我们需要的是“能解释行为且不依赖环境的最简
reward 表达 )
4.2 代码实现
专家经验中,获取(s,a,s')
1 2 3 4 5 6 7 8 9 10 11 E_states, E_actions, E_next_states = [], [], [] for traj in trajs: for (s, a, ns, r, d) in traj["transitions" ]: E_states.append(s) E_actions.append(a) E_next_states.append(ns) E_s = torch.FloatTensor(np.array(E_states)) E_a = F.one_hot(torch.LongTensor(E_actions), N_ACTIONS).float () E_ns = torch.FloatTensor(np.array(E_next_states)) n_expert = len (E_s)
初始化我们需要的网络
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 class PolicyNet (nn.Module): """策略网络: state → action logits""" def __init__ (self, state_dim=STATE_DIM, hidden=256 , n_actions=N_ACTIONS ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, n_actions) ) def forward (self, x ): return self .net(x) def act (self, state_np, greedy=False ): with torch.no_grad(): s = torch.FloatTensor(state_np).unsqueeze(0 ) logits = self .forward(s) if greedy: return logits.argmax(dim=-1 ).item() dist = Categorical(logits=logits) return dist.sample().item() class ValueNet (nn.Module): """价值网络: state → scalar (用于 PPO/IRL)""" def __init__ (self, state_dim=STATE_DIM, hidden=256 ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1 ) ) def forward (self, x ): return self .net(x).squeeze(-1 ) class AIRLRewardNet (nn.Module): """ AIRL g_θ 网络: (state, action) → scalar reward 对应论文中的可迁移奖励函数部分,仅依赖 (s, a) 而不依赖 s', 因此学到的奖励在不同环境动态下仍可迁移。 """ def __init__ (self, state_dim=STATE_DIM, n_actions=N_ACTIONS, hidden=128 ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim + n_actions, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1 ) ) def forward (self, states, actions_onehot ): """返回未经激活的 reward logit""" x = torch.cat([states, actions_onehot], dim=-1 ) return self .net(x).squeeze(-1 ) class AIRLShapingNet (nn.Module): """ AIRL h_φ 势函数网络: state → scalar 用于计算 reward shaping 项: γ·h(s') - h(s) 抵消对环境动态的依赖,使奖励函数可迁移。 """ def __init__ (self, state_dim=STATE_DIM, hidden=128 ): super ().__init__() self .net = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1 ) ) def forward (self, states ): return self .net(states).squeeze(-1 ) class AIRLDiscriminator (nn.Module): """ AIRL 完整判别器: f(s, a, s') = g_θ(s, a) + γ·h_φ(s') - h_φ(s) D(s, a, s') = sigmoid(f - log π(a|s)) 注意: forward 返回 f 的原始值 (logit),sigmoid 在训练损失中显式计算, 以便同时输出 reward (=f) 和 D (=sigmoid(f - log_pi))。 """ def __init__ (self, state_dim=STATE_DIM, n_actions=N_ACTIONS, hidden=128 , gamma=0.99 ): super ().__init__() self .gamma = gamma self .reward_net = AIRLRewardNet(state_dim, n_actions, hidden) self .shaping_net = AIRLShapingNet(state_dim, hidden) def f (self, states, actions_onehot, next_states ): """计算 f(s, a, s') = g(s,a) + γ·h(s') - h(s)""" g = self .reward_net(states, actions_onehot) h_s = self .shaping_net(states) h_s2 = self .shaping_net(next_states) return g + self .gamma * h_s2 - h_s def forward (self, states, actions_onehot, next_states, log_pi ): """ 返回判别器 logit: f(s,a,s') - log π(a|s) D = sigmoid(logit) ≈ P(expert | s,a,s') """ return self .f(states, actions_onehot, next_states) - log_pi def reward (self, states, actions_onehot, next_states ): """ 提取可迁移奖励 g_θ(s, a)(不含 shaping 项) 在迁移场景中使用此项;训练时使用 f 以保证稳定性。 """ return self .reward_net(states, actions_onehot) disc = AIRLDiscriminator(gamma=gamma) policy = PolicyNet() value = ValueNet() d_opt = Adam(disc.parameters(), lr=lr_disc, weight_decay=1e-5 ) p_opt = Adam(list (policy.parameters()) + list (value.parameters()), lr=lr_policy, weight_decay=1e-5 )
Policy和Value和之前一样。我们重点看disc网络。
其中 f函数计算的自然是我们的 \(f^∗(s,a,s')\) ,而
forward函数计算的则是 \(f-\log π\)
。
我们知道,收敛的时候 \(f^∗(s,a,s')=\logπ^∗(a|s)\) 因此 \(e^{f^*}=π^∗(a|s)\) 。我们得到 \[
f-\log π = \log e^{f}-\log π=\log{\frac{e^{f}}{π}}=\frac{π^*}{π}
\]
即当前动作越像专家,则数值越高。越像agent,则数值越低。我们可以直接把forward函数的输出当作奖励。
然后,如果我们将forward的输出直接sigmoid操作,就能还原成我们原本判别器的目标公式结果。
\[
D=σ(f-\log π) = \frac{\exp(f)}{\exp(f)+\pi(a|s)}
\] 具体推导的话,看这个
接下来就开始我们每轮的更新。
首先还是收集agent的数据。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 def _collect_rollout_airl (env, policy, disc, n_steps, gamma=0.99 ): """ AIRL rollout 收集。 reward = f(s,a,s') - log π(a|s) = g_θ(s,a) + γ·h_φ(s') - h_φ(s) - log π(a|s) 这是 AIRL 的内在 reward,不同于 GAIL: - 包含 next_state s'(用于 shaping 项) - 减去 log π(a|s) 使判别器更新与策略解耦 """ obs, _ = env.reset() rollout = [] disc.eval () for _ in range (n_steps): s = extract_state(obs, env) a = policy.act(s) next_obs, _, term, trunc, _ = env.step(a) done = term or trunc ns = extract_state(next_obs, env) with torch.no_grad(): s_t = torch.FloatTensor(s).unsqueeze(0 ) ns_t = torch.FloatTensor(ns).unsqueeze(0 ) a_oh = F.one_hot(torch.tensor([a]), N_ACTIONS).float () logits = policy(s_t) log_pi = Categorical(logits=logits).log_prob(torch.tensor([a])) f_val = disc(s_t, a_oh, ns_t, log_pi) r = f_val.item() rollout.append((s, a, ns, r, done)) obs = next_obs if done: obs, _ = env.reset() return rollout for ep in range (1 , airl_epochs + 1 ): rollout = _collect_rollout_airl(env, policy, disc, ppo_steps, gamma) A_s = torch.FloatTensor(np.array([t[0 ] for t in rollout])) A_a = F.one_hot(torch.LongTensor([t[1 ] for t in rollout]), N_ACTIONS).float () A_ns = torch.FloatTensor(np.array([t[2 ] for t in rollout]))
然后是更新判别器。
简单说就是从专家和agent经验中分别随机处2批batch_size的数据,然后计算出
d_expert = sigmoid(D(data_expert)) 和
a_expert = sigmoid(D(data_agent))
。然后带入我们判别器的loss目标中:
loss_d = ( F.binary_cross_entropy(d_expert,
torch.ones_like(d_expert)) + F.binary_cross_entropy(d_agent,
torch.zeros_like(d_agent)) )
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 for ep in range (1 , airl_epochs + 1 ): rollout = _collect_rollout_airl(env, policy, disc, ppo_steps, gamma) A_s = torch.FloatTensor(np.array([t[0 ] for t in rollout])) A_a = F.one_hot(torch.LongTensor([t[1 ] for t in rollout]), N_ACTIONS).float () A_ns = torch.FloatTensor(np.array([t[2 ] for t in rollout])) disc.train() last_disc_loss = 0.0 for _ in range (n_disc_updates): idx_e = torch.randperm(n_expert)[:min (batch_size, n_expert)] idx_a = torch.randperm(len (A_s))[:min (batch_size, len (A_s))] with torch.no_grad(): e_logits = policy(E_s[idx_e]) e_log_pi = Categorical(logits=e_logits).log_prob(torch.LongTensor(E_actions)[idx_e]) e_f = disc(E_s[idx_e], E_a[idx_e], E_ns[idx_e], e_log_pi.detach()) d_expert = torch.sigmoid(e_f) with torch.no_grad(): a_logits = policy(A_s[idx_a]) a_log_pi = Categorical(logits=a_logits).log_prob( torch.LongTensor([t[1 ] for t in rollout])[idx_a]) a_f = disc(A_s[idx_a], A_a[idx_a], A_ns[idx_a], a_log_pi.detach()) d_agent = torch.sigmoid(a_f) loss_d = ( F.binary_cross_entropy(d_expert, torch.ones_like(d_expert)) + F.binary_cross_entropy(d_agent, torch.zeros_like(d_agent)) ) d_opt.zero_grad() loss_d.backward() torch.nn.utils.clip_grad_norm_(disc.parameters(), 1.0 ) d_opt.step() last_disc_loss = loss_d.item()
然后是ppo的更新,和之前是一样的。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 def _ppo_update (policy, value, opt, rollout, epochs, gamma=0.99 , clip=0.2 , lam=0.95 ): """PPO clip 更新""" states = torch.FloatTensor(np.array([t[0 ] for t in rollout])) actions = torch.LongTensor([t[1 ] for t in rollout]) rewards = [t[3 ] for t in rollout] dones = [t[4 ] for t in rollout] values = value(states).detach().numpy() advs, returns = [], [] gae = 0.0 for i in reversed (range (len (rewards))): nv = values[i + 1 ] if i + 1 < len (values) else 0.0 mask = 0.0 if dones[i] else 1.0 delta = rewards[i] + gamma * nv * mask - values[i] gae = delta + gamma * lam * mask * gae advs.insert(0 , gae) returns.insert(0 , gae + values[i]) advs = torch.FloatTensor(advs) returns = torch.FloatTensor(returns) advs = (advs - advs.mean()) / (advs.std() + 1e-8 ) with torch.no_grad(): old_logp = Categorical(logits=policy(states)).log_prob(actions) N = len (states) for _ in range (epochs): idx = torch.randperm(N) for i in range (0 , N, 256 ): b = idx[i:i+256 ] logits = policy(states[b]) dist = Categorical(logits=logits) logp = dist.log_prob(actions[b]) ratio = (logp - old_logp[b]).exp() surr1 = ratio * advs[b] surr2 = ratio.clamp(1 -clip, 1 +clip) * advs[b] v_pred = value(states[b]) loss = (-torch.min (surr1, surr2).mean() + 0.5 * F.mse_loss(v_pred, returns[b]) - 0.01 * dist.entropy().mean()) opt.zero_grad(); loss.backward() torch.nn.utils.clip_grad_norm_( list (policy.parameters()) + list (value.parameters()), 0.5 ) opt.step() return loss.item() for ep in range (1 , airl_epochs + 1 ): rollout = _collect_rollout_airl(env, policy, disc, ppo_steps, gamma) A_s = torch.FloatTensor(np.array([t[0 ] for t in rollout])) A_a = F.one_hot(torch.LongTensor([t[1 ] for t in rollout]), N_ACTIONS).float () A_ns = torch.FloatTensor(np.array([t[2 ] for t in rollout])) disc.train() last_disc_loss = 0.0 for _ in range (n_disc_updates): idx_e = torch.randperm(n_expert)[:min (batch_size, n_expert)] idx_a = torch.randperm(len (A_s))[:min (batch_size, len (A_s))] with torch.no_grad(): e_logits = policy(E_s[idx_e]) e_log_pi = Categorical(logits=e_logits).log_prob(torch.LongTensor(E_actions)[idx_e]) e_f = disc(E_s[idx_e], E_a[idx_e], E_ns[idx_e], e_log_pi.detach()) d_expert = torch.sigmoid(e_f) with torch.no_grad(): a_logits = policy(A_s[idx_a]) a_log_pi = Categorical(logits=a_logits).log_prob( torch.LongTensor([t[1 ] for t in rollout])[idx_a]) a_f = disc(A_s[idx_a], A_a[idx_a], A_ns[idx_a], a_log_pi.detach()) d_agent = torch.sigmoid(a_f) loss_d = ( F.binary_cross_entropy(d_expert, torch.ones_like(d_expert)) + F.binary_cross_entropy(d_agent, torch.zeros_like(d_agent)) ) d_opt.zero_grad() loss_d.backward() torch.nn.utils.clip_grad_norm_(disc.parameters(), 1.0 ) d_opt.step() last_disc_loss = loss_d.item() disc_loss_hist.append(last_disc_loss) with torch.no_grad(): a_log_pi_all = Categorical(logits=policy(A_s)).log_prob( torch.LongTensor([t[1 ] for t in rollout])) f_vals = disc(A_s, A_a, A_ns, a_log_pi_all).detach() reward_f_hist.append(f_vals.mean().item()) disc.eval () _ppo_update(policy, value, p_opt, rollout, ppo_epochs)
5. AIRL的迁移学习
AIRL 学出来的奖励函数,相比于 IRL
更具解释性,因为它排除了环境动力学的影响。
IRL 在训练时隐式依赖了当前环境的转移动力学 \(p(s'∣s,a)\)
,奖励与特定环境耦合。比如同样是拿钥匙,在迷宫A里学到的奖励,换到迷宫B(地形不同,动力学不同)就失效了——IRL
学到的其实是"在这个环境的这个位置做这个动作有奖励 "。
AIRL 通过构造 \(f_\theta(s,a)\)
把动力学项约掉,学到的是"拿钥匙这个行为本身有奖励 ",与在哪个环境、怎么转移无关,因此可以迁移到不同环境,这才是解释性更强的真正含义。
所以,我们这里将在上一步5x5环境中学到的agent迁移到8x8的环境 中。
但毕竟是5x5的小环境到8x8的大环境,在小环境中基本上钥匙跟门都是在视野内的,大环境中则可能出现丢失的情况。因此agent也需要同步进行训练。
我们的方法是,agent在8x8环境中按照学好的策略收集数据。将成功的数据作为专家经验,重新训练判别器。同时,我们只要前30个步数最短的成功路径来当作我们的专家经验(希望学到最优策略)。
首先加载我们的模型。我们之前5x5环境训练好的模型保存为airl_policy.pt和airl_disc.pt,迁移训练的模型保存为transfer_policy.pt
和 transfer_disc.pt
。如果我们之前已经训练了一段时间有保存好的模型,我们就能直接加载继续训练,否则就读取5x5环境训练好的模型
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 AIRL_POLICY = DATA_DIR / "airl_policy.pt" AIRL_DISC = DATA_DIR / "airl_disc.pt" transfer_policy_path = DATA_DIR / "transfer_policy.pt" transfer_disc_path = DATA_DIR / "transfer_disc.pt" policy = PolicyNet() disc = AIRLDiscriminator(gamma=gamma) if resume and transfer_policy_path.exists() and transfer_disc_path.exists(): policy.load_state_dict( torch.load(transfer_policy_path, weights_only=True )) disc.load_state_dict( torch.load(transfer_disc_path, weights_only=True )) print ("续训:加载 transfer 权重" ) else : if not AIRL_POLICY.exists() or not AIRL_DISC.exists(): print ("找不到一阶段权重,请先运行 train_airl" ); return policy.load_state_dict( torch.load(AIRL_POLICY, weights_only=True )) disc.load_state_dict( torch.load(AIRL_DISC, weights_only=True )) print ("从一阶段 AIRL 权重开始迁移" ) value = ValueNet() p_opt = Adam(list (policy.parameters()) + list (value.parameters()), lr=lr_policy, weight_decay=1e-5 ) d_opt = Adam(disc.parameters(), lr=lr_disc, weight_decay=1e-5 )
然后加载专家数据,和之前airl一样,我们需要(s,a,s')。
让agent在新环境中执行20次,看下更新前的成功率。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 def _eval_policy (env, policy, n_eps=50 , greedy=True ): """跑 n_eps 局, 返回成功率""" successes = 0 for _ in range (n_eps): obs, _ = env.reset() for _ in range (200 ): s = extract_state(obs, env) a = policy.act(s, greedy=greedy) obs, reward, term, trunc, _ = env.step(a) if term or trunc: if reward > 0 : successes += 1 break return successes / n_eps print ("=== 自适应迁移训练 (8x8, disc 持续更新) ===\n" )baseline = _eval_policy(env, policy, 20 )
开始在新环境中训练,首先还是要agent收集数据。这次我们不仅收集数据,还要把成功的数据分出来,当作新环境的专家经验,因为在新环境我们不打算再手动制作经验。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 def _collect_rollout_with_successes (env, policy, disc, n_steps, gamma=0.99 ): """ 返回: rollout : [(s, a, ns, r_airl, done), ...] new_successes : [(len, is_5x5, s_list, a_list, ns_list), ...] """ obs, _ = env.reset() rollout = [] new_successes = [] current_episode = [] disc.eval () is_5x5 = (env.unwrapped.width == 5 ) for _ in range (n_steps): s = extract_state(obs, env) a = policy.act(s) next_obs, env_reward, term, trunc, _ = env.step(a) done = term or trunc ns = extract_state(next_obs, env) with torch.no_grad(): s_t = torch.FloatTensor(s).unsqueeze(0 ) ns_t = torch.FloatTensor(ns).unsqueeze(0 ) a_oh = F.one_hot(torch.tensor([a]), N_ACTIONS).float () logits = policy(s_t) log_pi = Categorical(logits=logits).log_prob(torch.tensor([a])) r = disc(s_t, a_oh, ns_t, log_pi).item() rollout.append((s, a, ns, r, done)) current_episode.append((s, a, ns)) if done: if env_reward > 0 : ep_len = len (current_episode) new_successes.append(( ep_len, is_5x5, [t[0 ] for t in current_episode], [t[1 ] for t in current_episode], [t[2 ] for t in current_episode], )) current_episode = [] obs, _ = env.reset() else : obs = next_obs return rollout, new_successes success_episodes = [] for ep in range (1 , ppo_epochs + 1 ): rollout, new_successes = _collect_rollout_with_successes( env, policy, disc, ppo_steps, gamma )
将获取的新专家经验放到经验池中。跟之前说的那样,我们只要最少步数的前30条经验。之前是有5x5旧的专家经验,因此写了先删除5x5经验,因为5x5的路径一般更短,后来改了但是懒得修改这一部分,反正不碍事。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 def update_success_pool (success_episodes, new_successes, max_size=50 ): """ success_episodes: [(len, is_5x5, s_list, a_list, ns_list), ...] 策略: 1. 总数 <= max_size 2. 优先删除 5x5 3. 否则删除最长 episode """ for ep in new_successes: success_episodes.append(ep) if len (success_episodes) > max_size: idx_5x5 = [i for i, e in enumerate (success_episodes) if e[1 ]] if idx_5x5: worst = max (idx_5x5, key=lambda i: success_episodes[i][0 ]) else : worst = max (range (len (success_episodes)), key=lambda i: success_episodes[i][0 ]) success_episodes.pop(worst) success_episodes = [] for ep in range (1 , ppo_epochs + 1 ): rollout, new_successes = _collect_rollout_with_successes( env, policy, disc, ppo_steps, gamma ) update_success_pool(success_episodes, new_successes, max_size=30 )
接下来就是跟之前一样的环节,获取专家数据
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 def flatten_success_episodes (success_episodes ): s_list, a_list, ns_list = [], [], [] for (_, _, s_seq, a_seq, ns_seq) in success_episodes: s_list.extend(s_seq) a_list.extend(a_seq) ns_list.extend(ns_seq) return s_list, a_list, ns_list success_episodes = [] for ep in range (1 , ppo_epochs + 1 ): rollout, new_successes = _collect_rollout_with_successes( env, policy, disc, ppo_steps, gamma ) update_success_pool(success_episodes, new_successes, max_size=30 ) expert_s, expert_a, expert_ns = flatten_success_episodes(success_episodes)
更新判别器
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 def _update_disc_adaptive (disc, d_opt, policy, expert_s, expert_a, expert_ns, agent_rollout, batch_size=256 , n_updates=1 ): """ 用成功缓冲区(作为专家)和当前 rollout(作为 agent)更新 disc。 expert_s/a/ns 接受 list 或 deque,统一转成 tensor。 """ E_s = torch.FloatTensor(np.array(list (expert_s))) E_a = F.one_hot(torch.LongTensor(list (expert_a)), N_ACTIONS).float () E_ns = torch.FloatTensor(np.array(list (expert_ns))) A_s = torch.FloatTensor(np.array([t[0 ] for t in agent_rollout])) A_a = F.one_hot(torch.LongTensor([t[1 ] for t in agent_rollout]), N_ACTIONS).float () A_ns = torch.FloatTensor(np.array([t[2 ] for t in agent_rollout])) A_actions = [t[1 ] for t in agent_rollout] disc.train() last_loss = 0.0 for _ in range (n_updates): idx_e = torch.randperm(len (E_s))[:min (batch_size, len (E_s))] idx_a = torch.randperm(len (A_s))[:min (batch_size, len (A_s))] with torch.no_grad(): e_log_pi = Categorical(logits=policy(E_s[idx_e])).log_prob( E_a[idx_e].argmax(dim=-1 )) a_log_pi = Categorical(logits=policy(A_s[idx_a])).log_prob( torch.LongTensor(A_actions)[idx_a]) e_f = disc(E_s[idx_e], E_a[idx_e], E_ns[idx_e], e_log_pi) a_f = disc(A_s[idx_a], A_a[idx_a], A_ns[idx_a], a_log_pi) loss = ( F.binary_cross_entropy(torch.sigmoid(e_f), torch.ones_like(e_f)) + F.binary_cross_entropy(torch.sigmoid(a_f), torch.zeros_like(a_f)) ) d_opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(disc.parameters(), 1.0 ) d_opt.step() last_loss = loss.item() disc.eval () return last_loss success_episodes = [] for ep in range (1 , ppo_epochs + 1 ): rollout, new_successes = _collect_rollout_with_successes( env, policy, disc, ppo_steps, gamma ) update_success_pool(success_episodes, new_successes, max_size=30 ) expert_s, expert_a, expert_ns = flatten_success_episodes(success_episodes) disc_loss = None if (len (expert_s) >= min_success_buffer and ep % disc_update_freq == 0 ): disc_loss = _update_disc_adaptive( disc, d_opt, policy, expert_s, expert_a, expert_ns, rollout, n_updates=1 , )
最后是PPO更新agent。(因为_ppo_update跟之前完全一样,我就不展示了)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 success_episodes = [] for ep in range (1 , ppo_epochs + 1 ): rollout, new_successes = _collect_rollout_with_successes( env, policy, disc, ppo_steps, gamma ) update_success_pool(success_episodes, new_successes, max_size=30 ) expert_s, expert_a, expert_ns = flatten_success_episodes(success_episodes) disc_loss = None if (len (expert_s) >= min_success_buffer and ep % disc_update_freq == 0 ): disc_loss = _update_disc_adaptive( disc, d_opt, policy, expert_s, expert_a, expert_ns, rollout, n_updates=1 , ) _ppo_update(policy, value, p_opt, rollout, epochs=5 )
最后添加保底机制。每更新5轮,就测试一次成功率,如果成功率最高,则保存新的模型。如果当前成功率比最高成功率低20%,则直接回滚成功率最高的模型。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 success_episodes = [] for ep in range (1 , ppo_epochs + 1 ): rollout, new_successes = _collect_rollout_with_successes( env, policy, disc, ppo_steps, gamma ) update_success_pool(success_episodes, new_successes, max_size=30 ) expert_s, expert_a, expert_ns = flatten_success_episodes(success_episodes) disc_loss = None if (len (expert_s) >= min_success_buffer and ep % disc_update_freq == 0 ): disc_loss = _update_disc_adaptive( disc, d_opt, policy, expert_s, expert_a, expert_ns, rollout, n_updates=1 , ) _ppo_update(policy, value, p_opt, rollout, epochs=5 ) rollback_thresh = 0.20 if ep % 5 == 0 : sr = _eval_policy(env, policy, 20 ) sr_hist.append((ep, sr)) if sr > best_sr: best_sr = sr torch.save(policy.state_dict(), transfer_policy_path) torch.save(disc.state_dict(), transfer_disc_path) flag = " ← 最优,已保存" elif best_sr > 0 and sr < best_sr - rollback_thresh: policy.load_state_dict( torch.load(transfer_policy_path, weights_only=True )) disc.load_state_dict( torch.load(transfer_disc_path, weights_only=True )) p_opt = Adam( list (policy.parameters()) + list (value.parameters()), lr=lr_policy, weight_decay=1e-5 ) d_opt = Adam(disc.parameters(), lr=lr_disc, weight_decay=1e-5 ) flag = " ← 下降过多,已回滚" else : flag = "" disc_str = f"{disc_loss:.4 f} " if disc_loss is not None else "跳过" print ( f" ep {ep:3d} /{ppo_epochs} | " f"sr={sr:.2 %} | best={best_sr:.2 %} | " f"disc_loss={disc_str} | " f"新成功={len (new_successes)} 条" f"{flag} " )
查看结果
修改奖励,期望最短步数。效果不怎么好
我们当前的仅用前30条步数最少的经验来训练,能有效减少无用步数
总结
BC 和 GAIL:学会怎么做 ,但不知道为什么
IRL 和 AIRL:推断为什么这样做 (恢复奖励函数)
AIRL 比 IRL
更进一步:恢复出的奖励函数可以迁移 到新环境