Skill Discovery | METRA：让策略探索 state 的紧凑 embedding space

论文标题：METRA: Scalable Unsupervised RL with Metric-Aware Abstraction
ICLR 2024 Oral。
arxiv：https://arxiv.org/abs/2310.08887
pdf：https://arxiv.org/pdf/2310.08887
html：https://arxiv.org/html/2310.08887
website：https://seohong.me/projects/metra/
GitHub：https://github.com/seohongpark/METRA

主要内容：

metra 关注 RL 的子领域 skill discovery，是目前 skill discovery 领域的 sota 方法之一。
skill discovery 希望以无监督的方式（没有 reward 信号）训练一系列策略 \(\pi(a|s,z)\) ，它们 condition on 一个高维向量 z，z 被称为 skill。总的来说，我们希望这一系列策略，可以彼此有很大的区分度，但又覆盖整个状态空间。具体的，这些策略满足两个特征：
- 1 可区分性：不同 skill z 生成的轨迹可以被区分，即，我们可以训练 / 得到一个 \(q(z|\tau)\)，来分辨一个轨迹是哪个 skill 生成的。
- 2 状态覆盖：这些策略应该尽可能地覆盖状态空间。
skill discovery 的通常做法：最大化 s 和 z 之间的互信息，\(I(s|z) = I(z|s) = H(s) - H(s|z) = H(z) - H(z|s)\) 。
metra 的核心思想：使用 temporal distance，为 state space 训练一个 embedding space，然后让 policy 在 embedding space 里跑的尽可能远。
- 训 embedding space 的动机：传统 skill discovery 的互信息方法，认为 ① 一个 humanoid 抬左脚 / 抬右脚 ② humanoid 往左跑 / 往右跑是没有区别的，认为这两个行为都是可分的。
- 然而，人类会觉得 ② 比 ① 区分度更大，这其实因为，抬左脚 / 抬右脚是很容易互相达到的两个状态（temporal distance 小），而 humanoid 在地图左边 / 右边，这两个状态不容易相互达到（temporal distance 大）。
- metra 发现，当我们跑一些复杂的环境，其中 state space 大一些，比如 state 维度高一些，传统 skill discovery 方法就很难覆盖很多 state 了。这是因为它们的很多 skill，都用来覆盖细枝末节、人类认为不重要的 state（temporal distance 小）。
- 因此，metra 认为，如果我们可以为 state space \(s\) 训练一个紧凑的 embedding space \(\phi(s)\)，使得 \(\phi(s_1),\phi(s_2)\) 之间的距离，与 \(s_1,s_2\) 之间的 temporal distance 相匹配；然后让 policy 尽可能覆盖 embedding space，便可以学到更有价值的 skill。

1 metra 的故事
2 metra 的理论 & 方法
3 metra 的实验 & 实验结果

1 metra 的故事

（deepseek 总结的，感觉比我写得好）

这篇论文讲的是如何在没有奖励信号的情况下，让 agent 自己摸索出各种有用的动作技能。这很重要，因为如果机器人能自己学会很多基础技能（比如走路、转身、拿东西），以后学具体任务（比如送快递）就会快得多。下面，分析它的动机和方法：

1.1 动机：现有方法的 Gap，METRA 如何解决

现有方法的 Gap：
- 方法一：瞎逛型（Pure Exploration）：让机器人尽量去没去过的地方。问题：在复杂环境（比如有很多关节的人形机器人）里，状态多到逛不完，效率太低。就像让你蒙着眼探索整个城市，可能一直绕圈子。
- 方法二：技能区分型（Mutual Info Skill）：让机器人学不同的技能，保证每个技能去的地方不一样。问题：它只关心技能“不一样”，不关心“多有用”。机器人可能学会“轻微抖腿”和“使劲抖腿”就算两种技能，但根本没移动！相当于厨师只会切不同形状的土豆丝，但不会炒菜。原文：“然而，它们存在一个共同局限，即它们往往最终发现的是简单、静态的行为，且状态覆盖范围有限”。
- 核心问题：这两种方法在简单环境还行，但在高维复杂环境（如像素输入的四足机器人、人形机器人）都搞不定，学不到真正有用的移动技能。
METRA 的动机，METRA 如何解决 Gap：
- 关键 Insight：不用费劲覆盖整个状态空间（太庞大），只需覆盖一个紧凑的“核心”空间（Latent Space Z），这个空间能代表环境中真正重要的变化方向。
- 如何连接“核心”空间和真实世界？：用时间距离（Temporal Distance） 当尺子。时间距离 = 两点间最快需要多少步到达。这很合理：在像素世界里，两个看起来不同的画面（比如机器人位置变化1米），如果一步就能到，那它们在这个“核心”空间里就该离得很近；如果需要跑 10 秒才能到，就该离得远。
- METRA 的目标：学一个映射函数 \(\phi(s)\) 把状态 \(s\)（比如像素图）映射到“核心”空间 \(Z\)（比如2维点），并满足：\(||\phi(s1) - \phi(s2)|| <= 时间距离(s1, s2)\)。同时，让机器人学会在这个空间里朝各个方向移动。
- 解决 Gap：既避开了“瞎逛”的低效，又避免了“技能区分”学无用技能的问题。专注于学习在最重要的维度（时间距离反映的） 上做长距离移动的技能。

1.2 方法：METRA 为什么合理和精妙

方法的核心：
- 两个一起学：
  1. 映射函数 \(\phi(s)\)：把状态（如像素图）映射到低维空间 \(Z\)，并遵守 \(||\phi(s) - \phi(s')|| <= 1\)（如果 \(s\) 和 \(s'\) 是相邻状态）。这保证了 \(Z\) 空间的距离能反映时间距离。
  2. 技能策略 \(π(a|s, z)\)：目标是让机器人执行技能 \(z\) 时，在 \(Z\) 空间里尽量沿着 \(z\) 指定的方向移动。
- 奖励函数：\(r(s, z, s') = (\phi(s') - \phi(s))^T z\)。
  - 解释：如果机器人成功让状态在 \(Z\) 空间移动了 \(Δ\phi = \phi(s') - \phi(s)\)，并且这个移动方向 \(Δ\phi\) 和它指定的方向 \(z\) 很一致（点积大），就给它高奖励。
  - Intuition：这等于告诉机器人：“你选一个方向 \(z\)，然后努力让状态在 \(Z\) 空间里沿着这个方向走得越远越好”。因为 \(Z\) 空间被约束得能反映真实的时间距离，在这里走得远，相当于在真实世界完成了有意义的移动（比如从房间一头走到另一头）。
为什么合理 & 精妙？
1. 时间距离是普适的“好尺子”：它只依赖于环境本身的动态（几步能到），不依赖于状态的具体表示（比如像素值），完美适用于高维像素输入环境。
2. 约束 \(||Δ\phi|| <= 1\) 是安全的保证：它确保 \(Z\) 空间不会扭曲现实。你在 \(Z\) 空间移动 1 单位，真实世界至少需要 1 步，学到的技能反映真实的可达性。
3. 最大化点积驱动多样性探索：为了最大化奖励，机器人会倾向于：
  - 学那些能在 \(Z\) 空间产生大位移 \(Δ\phi\) 的技能（移动得远）。
  - 让不同 \(z\) 指向 \(Z\) 空间中不同的方向（覆盖不同方向）。
  - 这恰好让机器人学会了在最重要的维度上做长距离移动的各种技能。
4. 自洽的闭环：\(\phi\) 约束 \(Z\) 空间的结构，\(π\) 利用这个结构学习移动技能，技能收集到的数据又反过来优化 \(\phi\)。两者相互促进。

1.3 总结一下这个故事

动机故事：以前的无监督 RL 在复杂环境（尤其是像素）不是逛得太低效（瞎逛），就是学些没用的花拳绣腿（互信息技能）。根本原因是它们没抓住重点（环境的关键变化维度）。
解决方法故事：我们提出 METRA，用时间距离这把万能钥匙，构建一个紧凑的“核心地图”（\(Z\) 空间）。机器人只需学习在这个地图上朝各个方向快速移动的技能。这样既能高效探索，又能学到真正有意义（反映长距离移动）的技能。这把钥匙（时间距离）让我们的方法天然适应像素环境，约束保证了地图不骗人，奖励机制驱动了有效探索。
结果：METRA 首次在像素输入的四足机器人（Quadruped）和人形机器人（Humanoid）上学会了多样的移动技能，并且在多个任务上超越了之前的 SOTA 方法。

2 metra 的理论 & 方法

metra 所做的第一件事，就是把 skill discovery 里经典的互信息目标，换成了 Wasserstein dependency measure（WDM）。关于 WDM 的讲解，可以看本站博客，关于把互信息换成 WDM 动机，本站博客里也有写。总之，我们的互信息 \(I(s,z)\) 现在变成了这样：\(I_W(s,z)\)。

根据本站博客，可以这样计算 \(I_W(s,z)\) ：

\[I_W(S;Z) = \sup_{\|f\|_{\mathcal L}\le 1}\,
\Bigl\{\,
\mathbb E_{p(s,z)}\bigl[f(s,z)\bigr]
-
\mathbb E_{p(s)\,p(z)}\bigl[f(s,z)\bigr]
\Bigr\} \tag{1}
\]

这里的 \(f(s,z)\) 相当于 WDM 里所讲的相似性函数，它应该满足 1-Lipschitz 条件 \(\|f\|_{\mathcal L}\le 1\)，即，\(|f(s_1,z_1) - f(s_2,z_2) | \le \|(s_1,z_1) - (s_2,z_2)\|_d\) ，其中 \(\|\cdot \|_d\) 是任选的可以度量向量之间距离的 metric，下文 metra 把它选成了 temporal distance。

metra 声称，(1) 式的形式已经可以使用 WDM 的方法来训，同时，我们可以训一个策略 \(\pi(s|s,z)\)，它的 reward 是 \(r(s,z) = f(s,z) -(1/N)\sum_{i=1}^N f(s,z_i)\) 。然而，这个形式还可以再简化。

metra 将 \(f(s,z)\) 设为 \(\phi(s)^Tz\) 的形式，把 (1) 式转换为以下形式：

\[I_W(S;Z) \approx \sup_{\|\phi\|_{\mathcal L}\le 1}\,
\Bigl\{\,
\mathbb E_{p(s,z)}\bigl[\phi(s)^Tz\bigr]
-
\mathbb E_{p(s)}\big[\phi(s)\big]^T \mathbb E_{p(z)}\big[z\big]
\Bigr\} \tag{2}
\]

这里是约等于，因为其实 LHS RHS 并不完全相等。metra 声称，\(\|\phi\|_{\mathcal L}\le 1\)（即 \(\|\phi(s_1) - \phi(s_2) \|_2 \le \|s_1 - s_2\|_d\)）并不等价于 \(\|f\|_{\mathcal L}\le 1\)，但反正 \(\|\phi\|_{\mathcal L}\le 1\) 更好算。并且，\(\|f\|_{\mathcal L}\) 的 upper bound 可以用 \(\|\phi\|_{\mathcal L}\) 等一堆东西表示出来，所以 metra 认为这样做是完全可以的。我不懂数学，不过感觉这个形式挺好看的。

（metra 原文是这样写的：

\[I_W(S;Z) \approx \sup_{\|\phi\|_{\mathcal L}\le 1 ,\, \|\psi\|_{\mathcal L}\le 1}\,
\Bigl\{\,
\mathbb E_{p(s,z)}\bigl[\phi(s)^T \psi(z)\bigr]
-
\mathbb E_{p(s)}\big[\phi(s)\big]^T \mathbb E_{p(z)}\big[\psi(z)\big]
\Bigr\} \tag{3}
\]

然后在后文，把 \(\psi(z)\) 设为 \(z\)。）

接下来，metra 继续变换 (2) 式。metra 让 \(I_W(S;Z)\) 只考虑一个 trajectory 里的最后一个状态，即 \(I_W(S_T;Z)\)。然后，metra 进行一个裂项相消，得到：

\[I_W(S_T;Z) \approx \sup_{\|\phi\|_L \leq 1} \mathbb{E}_{p(\tau, z)}[\phi(S_T)^\top z] - \mathbb{E}_{p(\tau)}[\phi(S_T)]^\top \mathbb{E}_{p(z)}[z] \\
= \sup_{\phi} \sum_{t=0}^{T-1} \Big( \mathbb{E}_{p(\tau, z)}[(\phi(s_{t+1}) - \phi(s_t))^\top z] - \mathbb{E}_{p(\tau)}[\phi(s_{t+1}) - \phi(s_t)]^\top \mathbb{E}_{p(z)}[z] \Big) \tag{4}
\]

这样的话，还应该剩一个 \(\mathbb{E}_{p(\tau, z)}[\phi(s_0)^\top z] - \mathbb{E}_{p(\tau)}[\phi(s_{0})]^\top \mathbb{E}_{p(z)}[z]\) 才对。不过， \(s_0, z\) 是相互独立的，这一项就 = 0 了。

我们把 (4) 式整理一下，就变成了：

\[I_W(S_T;Z) \approx \sup_{\|\phi\|_L \leq 1} \mathbb{E}_{p(\tau, z)} \bigg[\sum_{t=0}^{T-1} \big[\phi(s_{t+1}) - \phi(s_{t})\big]^T (z-\bar z) \bigg] \tag{5}
\]

然后，我们假设 skill 的平均值 \(\bar z=0\)，这很容易做到，因为在训练过程中，z 的值其实是我们采样得到的。

现在，Lipschitz 条件 \(\|\phi\|_L \leq 1\) 还不知道该怎么处理。这里的关键在于，我们要去优化 \(\|\phi(s_1) - \phi(s_2) \|_2 \le \|s_1 - s_2\|_d\) 的目标，但是 metric d 还不知道该选什么。

metra 选择 metric d 为 temporal distance（时序距离，可参考本站博客）。temporal distance 的大致含义，是我们去训练一个 state embedding \(\phi(s)\) ，使得容易相互到达的 state 的 embedding 离得近。具体的，

\[\|\phi(s_{t+1}) - \phi(s_t) \|_2 \le 1 \tag{6}
\]

我们发现，这跟 Lipschitz 条件的形式很相像，就直接用 \(\|\phi(s_{t+1}) - \phi(s_t) \|_2 \le 1\) 来作为 (5) 式的约束了。现在，(5) 式变成了这样：

\[\sup_{\pi, \phi} \quad \mathbb{E}_{p(\tau, z)} \left[ \sum_{t=0}^{T-1} (\phi(s_{t+1}) - \phi(s_t))^\top z \right] \quad \text{s.t.} \quad \|\phi(s) - \phi(s')\|_2 \leq 1 \tag{7}
\]

具体的，上式可以被拆为三部分，进行跟 WDM 差不多的对偶 learning：

训练 state embedding \(\phi(s)\) ：\(\max \mathbb{E}_{(s,z,s') \sim \mathcal{D}} [(\phi(s') - \phi(s))^\top z + \lambda \cdot \min(\varepsilon, 1 - \|\phi(s) - \phi(s')\|_2^2)]\) ；
更新 Lagrange 乘子 \(\lambda\) ：\(\min \mathbb{E}_{(s,z,s') \sim \mathcal{D}} [\lambda \cdot \min(\varepsilon, 1 - \|\phi(s) - \phi(s')\|_2^2)]\) ；
更新 policy \(\pi(a|s,z)\) ：reward \(r(s, z, s') = (\phi(s') - \phi(s))^\top z\)，使用 SAC 算法。
从 reward 的角度来看，突然感觉 ETD 跟 metra 真像（

3 metra 的实验 & 实验结果

3.1 实验环境

state-based 环境：
- Ant（29 维状态）：四足机器人运动。
- HalfCheetah（18 维状态）：仿猎豹机器人奔跑。
pixel-based 环境：
- Quadruped & Humanoid（64×64×3 图像）：四足 / 人形机器人运动。
- Kitchen（64×64×3 图像）：机器人完成 6 项厨房子任务，烧水 / 微波炉 / 开关灯等。

3.2 baseline

skill discovery 方法：

LSD：最大化状态间欧氏距离，还没读过，貌似是 metra 作者之前的工作。
DIAYN：互信息最大化区分技能，是 Eysenbach 的工作，skill discovery 最早的工作之一。
DADS：学习技能动力学模型，还没读过。
CIC：对比学习 + 探索奖励混合，还没读过。
APS：使用 successor feature 的 skill discovery 工作，有点忘记具体内容了。

探索方法：
- ICM / RND / APT：基于状态新颖性的内在奖励。
- LBS：基于贝叶斯惊喜的探索，没读过。
- Plan2Explore：动态模型不确定性最大化，没读过。
目标达成方法：
- LEXA：联合训练探索策略 + 目标条件策略，没读过。

3.3 评估指标

策略状态覆盖率：评估当前策略在 状态空间 的覆盖范围。并不考察整个 state space 的覆盖率，而是 state space 的一个子集，如 Ant / Humanoid 记录 x-y 平面被访问的 1×1 方格数。
任务覆盖率：对于 Kitchen 环境，统计 6 项子任务 完成数量。
下游任务性能：使用 分层控制，冻结 low-level 的 skill-conditioned policy，训练一个 high-level 策略，优化具体的任务奖励。
零样本目标达成：用学习到的技能 直接执行 目标条件任务，如 Humanoid 中计算与目标的欧氏距离。

3.4 实验结果

技能多样性（定性的）：
- metra 是 首个在像素 Quadruped / Humanoid 中发现运动技能 的方法。LSD / DIAYN 等仅在状态环境有效，像素环境完全失败。
- （metra 把 pixel-based 环境的地面染成了五彩的颜色，往每个方向跑，都有不同的颜色。我一开始以为，这是为了可视化好看，后来意识到，只有这样，state embedding \(\phi(s)\) 才能分辨出来往不同方向跑的小蜘蛛，彼此之间有区别。）
状态覆盖率（以下都是定量的）：
- metra 在 Ant 覆盖率超 2000，比第二名 LSD（~500）高4倍（图5）
- 在 Humanoid 覆盖率 75，远超探索方法（<25）（图8）
下游任务性能：
- metra 在5个任务中4项排名第一，例如：
  - Quadruped Goal 任务得分 6.0，LSD 仅 2.5。
  - HalfCheetah Hurdle 跳跃任务得分 7.5，DADS 仅 5.0。
零样本目标达成：
- metra 在 Ant 目标距离减少 40 单位，LEXA 仅 20。
- Kitchen 任务完成数 4/6，LEXA 为 3/6。