强化学习 11. 模仿学习(原理与实验)

前言

从零开始学习ai文章系列计划是个人在《动手学深度学习》和《磨菇书》两本书的学习中的个人笔记,文章也会以课本中的章节分开,即每个章节一片笔记。我会尽量的把主要内容以及遇到的难点进行记录与解决,如果哪里有错误的欢迎指正。或者不清晰的可以直接查看原文部分。

《蘑菇书》原文(课本):https://datawhalechina.github.io/easy-rl/#/

(由于有时候公式太多,可能会直接贴图片)

蘑菇书的文章结构不会跟之前《动手学深度学习》按照原文章节进行,个人会适当调节。

(第十章实验鸽了,好奇心实验做的道心破碎,本章是蘑菇书学习最后一章)


0. 实验环境

实验在 MiniGrid-DoorKey-8x8-v0MiniGrid-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:特殊地面对象(有些环境才用)

动作空间

终止条件:

  1. 捡起钥匙,打开门后,进入到绿色方块处。
  2. 超过最大步数。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"
# ─── 动作映射 ─────────────────────────────────────────────────────────────────
# MiniGrid 动作: 0=left, 1=right, 2=forward, 3=pickup, 4=drop, 5=toggle, 6=done
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 = [] # [(state, action, next_state, reward, done), ...]
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:
# 自动判断: 前方是门用 toggle, 前方是钥匙用 pickup
_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:.2f}"
f" {'✓ 成功! 按 S 保存,或按 R 丢弃重来' if success else '✗ 超时,按 R 重置'}")
# 不自动重置,等待玩家手动按 S 保存 或 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,      # [0]   朝向
carrying, # [1] 是否持有钥匙
key_relative, # [2,3] 钥匙相对位置
door_relative, # [4,5] 门相对位置
door_open, # [6] 门是否已开
dist_to_key, # [7] 到钥匙的距离

实际上我们不应该结合世界坐标,因为这是未知的,但为了训练的方便我们也放进去了。

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_θ 网络在迁移时
收到训练分布外的输入,输出混乱的奖励信号。
相对坐标在任意地图尺寸下语义一致。
"""
# ── 局部视野图像 (7x7x3),归一化 ──────────────────────────────────────
image = obs["image"].astype(np.float32) / 10.0 # (7,7,3)
image_flat = image.flatten() # 147 维

# ── 穿透 gymnasium 包装器,获取底层 MiniGridEnv ───────────────────────
inner = env
while hasattr(inner, "env"):
inner = inner.env

width = float(inner.width) # 用于归一化,5x5→5,8x8→8

# ── 基础标量 ──────────────────────────────────────────────────────────
agent_pos = np.array(inner.agent_pos, dtype=np.float32) # 绝对坐标,仅内部用
agent_dir = np.array([inner.agent_dir / 3.0], dtype=np.float32) # [0,1]
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)

# ── 计算相对坐标(核心修改)──────────────────────────────────────────
#
# 相对坐标 = (目标绝对坐标 - agent绝对坐标) / width
# 范围:(-1, 1),与地图尺寸无关
#
# 持有钥匙时,钥匙已不在地图上,key_pos_abs 扫不到 → 置 (0,0)
# 这同时也告诉网络"钥匙问题已解决"

if key_pos_abs is not None:
key_relative = (key_pos_abs - agent_pos) / width # (dx, dy)
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 # (dx, dy)
else:
door_relative = np.array([0.0, 0.0], dtype=np.float32)

# ── 辅助距离特征 ─────────────────────────────────────────────────────
# 曼哈顿距离归一化,显式告诉网络"远近"关系
# 持有钥匙时为 0(已完成靠近目标)
if key_pos_abs is not None:
dist_to_key = np.array(
[np.abs(key_relative).sum()], # |dx|+|dy|,已归一化
dtype=np.float32
)
else:
dist_to_key = np.array([0.0], dtype=np.float32)

# ── 拼接 ─────────────────────────────────────────────────────────────
# agent_dir(1) + carrying(1) + key_relative(2) +
# door_relative(2) + door_open(1) + dist_to_key(1) = 8 维
scalars = np.concatenate([
agent_dir, # [0] 朝向
carrying, # [1] 是否持有钥匙
key_relative, # [2,3] 钥匙相对位置
door_relative, # [4,5] 门相对位置
door_open, # [6] 门是否已开
dist_to_key, # [7] 到钥匙的距离
]) # 共 8 维

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):
# 1.读取专家数据文件
if not TRAJ_FILE.exists():
print("无轨迹文件, 请先 collect");
return
with open(TRAJ_FILE, "rb") as f:
trajs = pickle.load(f)
# 构建数据集,只需要其中的s,a
states, actions = [], []
for traj in trajs:
for (s, a, ns, r, d) in traj["transitions"]:
states.append(s)
actions.append(a)
# 转成tensor
states = torch.FloatTensor(np.array(states))
actions = torch.LongTensor(actions)
N = len(states)
print(f"BC 数据: {N} 条 (s,a) 对, 来自 {len(trajs)} 条轨迹")

# 2. 初始化policy网络
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")


# 3.开始训练
for epoch in range(1, epochs + 1):
idx = torch.randperm(N)
total_loss = 0.0
# 每轮从专家数据中随机拿出batch_size个s,a
# 计算交叉熵 CE(π(s), a_expert)
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:.4f} | "
f"train_acc={acc:.3f} | lr={sched.get_last_lr()[0]:.5f}")

_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)] \] 然后提取常数项,log⁡Z(θ) 与具体哪条轨迹 τ 无关,对 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 的参数。

此时:

  • 特征提取
  • 特征组合
  • reward 映射

都已经被神经网络隐式学习,因此不再显式构造 \(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):
# ── 内循环: PPO 收集数据 ──────────────────────────────────────────
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):
# ── 内循环: PPO 收集数据 ──────────────────────────────────────────
rollout = _collect_rollout(env, policy, reward_net, ppo_steps)
agent_states = torch.FloatTensor(np.array([t[0] for t in rollout]))

# ── 更新 Reward Net (MaxEnt IRL) ─────────────────────────────────
# 只在 policy 还没很好时更新 reward(成功率<90% 才更新)
# 避免好 policy 被乱掉的 reward 拖垮
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]

# GAE 计算 advantage
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)

# 旧策略 log prob
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):
# ── 内循环: PPO 收集数据 ──────────────────────────────────────────
rollout = _collect_rollout(env, policy, reward_net, ppo_steps)
agent_states = torch.FloatTensor(np.array([t[0] for t in rollout]))

# ── 更新 Reward Net (MaxEnt IRL) ─────────────────────────────────
# 只在 policy 还没很好时更新 reward(成功率<90% 才更新)
# 避免好 policy 被乱掉的 reward 拖垮
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 更新 policy ───────────────────────────────────────
_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
# 旧策略 log prob
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) # log-odds as reward
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):
# ── 收集智能体轨迹 (使用判别器 reward) ────────────────────────────
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()

# ── 更新判别器 D ───────────────────────────────────────────────────
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])
# BCE: 专家标签=1, 智能体标签=0
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]

# GAE 计算 advantage
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)

# 旧策略 log prob
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):
# ── 收集智能体轨迹 (使用判别器 reward) ────────────────────────────
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()

# ── 更新判别器 D ───────────────────────────────────────────────────
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])
# BCE: 专家标签=1, 智能体标签=0
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 更新 policy ────────────────────────────────────────────────
_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():.4f} | 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) # (B,)
h_s = self.shaping_net(states) # (B,)
h_s2 = self.shaping_net(next_states) # (B,)
return g + self.gamma * h_s2 - h_s # (B,)

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()

# log π(a|s)
logits = policy(s_t)
log_pi = Categorical(logits=logits).log_prob(torch.tensor([a]))

# f(s,a,s') - log π(a|s) ← AIRL reward
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):

# ════════════════════════════════════════════════════════════════════
# Step 1: 用当前策略收集 rollout,计算 AIRL reward
# ════════════════════════════════════════════════════════════════════
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):

# ════════════════════════════════════════════════════════════════════
# Step 1: 用当前策略收集 rollout,计算 AIRL reward
# ════════════════════════════════════════════════════════════════════
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]))

# ════════════════════════════════════════════════════════════════════
# Step 2: 更新 AIRL 判别器
# 损失: BCE(σ(f - log π), 专家=1 / 智能体=0)
# 等价于最大化专家的 f(s,a,s') - log π(a|s),
# 最小化智能体的 f(s,a,s') - log π(a|s)
# ════════════════════════════════════════════════════════════════════
disc.train()
last_disc_loss = 0.0

for _ in range(n_disc_updates):
# 随机采样 mini-batch
idx_e = torch.randperm(n_expert)[:min(batch_size, n_expert)]
idx_a = torch.randperm(len(A_s))[:min(batch_size, len(A_s))]

# ── 专家样本的判别器 logit ────────────────────────────────────
with torch.no_grad():
# 专家动作的 log π(a|s)(来自当前策略)
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)

# ── 智能体样本的判别器 logit ──────────────────────────────────
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)

# ── BCE 损失 ─────────────────────────────────────────────────
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]

# GAE 计算 advantage
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)

# 旧策略 log prob
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):

# ════════════════════════════════════════════════════════════════════
# Step 1: 用当前策略收集 rollout,计算 AIRL reward
# ════════════════════════════════════════════════════════════════════
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]))

# ════════════════════════════════════════════════════════════════════
# Step 2: 更新 AIRL 判别器
# 损失: BCE(σ(f - log π), 专家=1 / 智能体=0)
# 等价于最大化专家的 f(s,a,s') - log π(a|s),
# 最小化智能体的 f(s,a,s') - log π(a|s)
# ════════════════════════════════════════════════════════════════════
disc.train()
last_disc_loss = 0.0

for _ in range(n_disc_updates):
# 随机采样 mini-batch
idx_e = torch.randperm(n_expert)[:min(batch_size, n_expert)]
idx_a = torch.randperm(len(A_s))[:min(batch_size, len(A_s))]

# ── 专家样本的判别器 logit ────────────────────────────────────
with torch.no_grad():
# 专家动作的 log π(a|s)(来自当前策略)
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 = sigmoid(f_expert)
d_expert = torch.sigmoid(e_f)

# ── 智能体样本的判别器 logit ──────────────────────────────────
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)

# ── BCE 损失 ─────────────────────────────────────────────────
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)

# 记录当前 rollout 的平均 f 值(衡量奖励函数的区分力)
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())

# ════════════════════════════════════════════════════════════════════
# Step 3: PPO 更新策略(使用 AIRL reward)
# ════════════════════════════════════════════════════════════════════
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.ptairl_disc.pt,迁移训练的模型保存为transfer_policy.pttransfer_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()

# 判断当前是不是 5x5
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):

# ── Step 1: rollout ─────────────────────────────
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:

# 找所有 5x5
idx_5x5 = [i for i, e in enumerate(success_episodes) if e[1]]

if idx_5x5:
# 删除最长的 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):

# ── Step 1: rollout ─────────────────────────────
rollout, new_successes = _collect_rollout_with_successes(
env, policy, disc, ppo_steps, gamma
)

# ── Step 2: 更新“最优成功池” ────────────────────
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):

# ── Step 1: rollout ─────────────────────────────
rollout, new_successes = _collect_rollout_with_successes(
env, policy, disc, ppo_steps, gamma
)

# ── Step 2: 更新“最优成功池” ────────────────────
update_success_pool(success_episodes, new_successes, max_size=30)

# 展开成 step-level 数据给 disc
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): # 默认只更新 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):

# ── Step 1: rollout ─────────────────────────────
rollout, new_successes = _collect_rollout_with_successes(
env, policy, disc, ppo_steps, gamma
)

# ── Step 2: 更新“最优成功池” ────────────────────
update_success_pool(success_episodes, new_successes, max_size=30)

# 展开成 step-level 数据给 disc
expert_s, expert_a, expert_ns = flatten_success_episodes(success_episodes)

# ── Step 3: 更新 disc ──────────────────────────
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):

# ── Step 1: rollout ─────────────────────────────
rollout, new_successes = _collect_rollout_with_successes(
env, policy, disc, ppo_steps, gamma
)

# ── Step 2: 更新“最优成功池” ────────────────────
update_success_pool(success_episodes, new_successes, max_size=30)

# 展开成 step-level 数据给 disc
expert_s, expert_a, expert_ns = flatten_success_episodes(success_episodes)

# ── Step 3: 更新 disc ──────────────────────────
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,
)

# ── Step 4: PPO ────────────────────────────────
_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):

# ── Step 1: rollout ─────────────────────────────
rollout, new_successes = _collect_rollout_with_successes(
env, policy, disc, ppo_steps, gamma
)

# ── Step 2: 更新“最优成功池” ────────────────────
update_success_pool(success_episodes, new_successes, max_size=30)

# 展开成 step-level 数据给 disc
expert_s, expert_a, expert_ns = flatten_success_episodes(success_episodes)

# ── Step 3: 更新 disc ──────────────────────────
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,
)

# ── Step 4: PPO ────────────────────────────────
_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:
# 新高:同步保存 policy 和 disc,保证两者始终匹配
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 + disc + 优化器
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:.4f}" 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"buf={len(success_buf_s)}条 | "
f"新成功={len(new_successes)}条"
f"{flag}"
)

查看结果

修改奖励,期望最短步数。效果不怎么好

image-20260504192732297

我们当前的仅用前30条步数最少的经验来训练,能有效减少无用步数

总结

BC 和 GAIL:学会怎么做,但不知道为什么

IRL 和 AIRL:推断为什么这样做(恢复奖励函数)

AIRL 比 IRL 更进一步:恢复出的奖励函数可以迁移到新环境