結構化策略初始化加速大規模離散動作空間離線強化學習

IMPROVING AND ACCELERATING OFFLINE RL INLARGE DISCRETE ACTION SPACES WITH STRUCTUREDPOLICY INITIALIZATION

結構化策略初始化加速大規模離散動作空間離線強化學習

https://arxiv.org/pdf/2601.04441

摘要

在離散組合動作空間中進行強化學習，需搜索指數級數量的聯合動作，以同時選擇多個能形成協調組合的子動作。現有方法要么通過假設子動作相互獨立來簡化策略學習，但這常導致不協調或無效的動作；要么嘗試聯合學習動作結構與控制，但這往往緩慢且不穩定。我們提出了結構化策略初始化（SPIN），這是一種兩階段框架：首先預訓練動作結構模型（ASM）以捕捉有效動作的流形，隨后凍結該表示，并訓練輕量級策略頭用于控制。在具有挑戰性的離散DM Control基準測試中，SPIN較當前最優方法將平均回報最高提升了39%，同時將收斂時間最多縮短至原來的1/12.8（即提速12.8倍）。

1 引言

許多現實世界的問題需要在高維離散動作空間中進行決策，包括醫療保健（Liu et al., 2020）、機器人裝配（Driess et al., 2020）、推薦系統（Zhao et al., 2018）以及網約車（Lin et al., 2018）等領域的應用。在這些領域中，在線探索（online exploration）可能代價高昂或不安全，這使得離線強化學習（RL）（Lange et al., 2012; Levine et al., 2020）成為一個極具吸引力的框架。然而，標準的離線RL方法（Fujimoto et al., 2019; Agarwal et al., 2020; Fu et al., 2020; Kumar et al., 2020; Kostrikov et al., 2021）并非專為大型離散動作空間設計，因為它們需要在整個離散動作集上最大化Q函數或對策略進行參數化——隨著動作空間隨
呈指數級擴展，這些操作將變得難以處理（intractable），其中 A 是子動作維度的數量，是每個維度的可選數量。

在這些復雜場景中進行學習需要解決兩個相關問題：（i）在指數級數量的聯合動作中進行搜索，以及（ii）確保所選子動作形成協調一致的組合。針對此類組合空間設計的方法傳統上通過施加強結構先驗（例如假設子動作間條件獨立）來簡化策略學習（Tang et al., 2022; Beeson et al., 2024）。然而，這犧牲了表示能力，致使模型無法捕捉有效控制所需的交互作用。其他方法嘗試同時學習動作表示與優化策略（Zhang et al., 2018; Landers et al., 2024; 2025），但這種目標耦合往往導致學習過程緩慢且不穩定。

我們提出了結構化策略初始化（Structured Policy Initialization, SPIN），這是一個將表示學習與控制解耦的兩階段框架。在第一階段，通過自監督訓練一個動作結構模型（Action Structure Model, ASM），以學習一個表示函數；該函數以狀態 s s 為條件，在子動作上誘導產生一個特征空間，其中結構上連貫的聯合動作集中在一個低維流形上。隨后，在第二階段凍結該動作空間表示，此時控制問題簡化為針對下游強化學習（RL）任務，在動作流形上學習輕量級策略頭。通過先學習結構再學習策略，SPIN 允許智能體利用底層的動作幾何結構，而不是在原始組合空間中進行搜索。這帶來了更快的訓練速度和提升的策略性能（圖 1）。在數據集大小和質量、動作維度以及動作基數各異的多樣化基準測試中，SPIN 相比當前最先進方法（state of the art）將平均回報最高提升了 39%，并將達到最先進性能所需的訓練時間最多縮短了 12.8 倍。

我們的貢獻如下：

• 我們將離散結構化動作空間中的離線強化學習（RL）重新構建為一個表示問題，將動作結構學習與控制分離開來。
• 我們提出了 SPIN，這是一個兩階段框架，通過預訓練并凍結動作空間表示來加速并改進策略學習。
• 我們展示了 SPIN 在具有挑戰性的基準測試中實現了最先進（SOTA）的性能，在顯著更快的同時優于現有方法。
• 我們分析了學習到的表示，以證明在離散組合動作空間中進行有效策略學習時，捕捉動作結構至關重要。

2 相關工作

大離散動作空間中的強化學習。 針對路由（Nazari et al., 2018; Delarue et al., 2020）和資源分配（Chen et al., 2024）等領域的組合動作空間，已開發出多種強化學習方法，但這些方法通常依賴于特定任務的知識。研究也引入了通用方法（Dulac-Arnold et al., 2015; Tavakoli et al., 2018; Farquhar et al., 2020; Van de Wiele et al., 2020; Zhao et al., 2023），但它們通常面向在線學習設計，難以直接適應離線數據集的約束條件。在離線強化學習中，現有方法通常對策略或 Q 函數進行分解（factorize）（Tang et al., 2022; Beeson et al., 2024）。然而，這種分解強制子動作之間滿足條件獨立性，從而限制了模型的表征能力，并在子動作存在強依賴關系時失效。其他方法則顯式地捕捉依賴關系——例如 BraVE（Landers et al., 2024）對跨維度交互進行建模，但其計算復雜度隨動作規模擴大而急劇增加（擴展性差）；而自回歸策略（Zhang et al., 2018）則強加了固定的動作順序，破壞了排列不變性。最近，SAINT（Landers et al., 2025）引入了一種基于 Transformer 的策略，通過自注意力機制捕捉子動作間的依賴關系，但其聯合學習動作結構與控制策略的方式，導致了訓練緩慢且不穩定。另一條相關研究線致力于為大規模但平坦（flat）的動作空間學習表征。其中最相關的是 MERLION（Gu et al., 2022），它為離線強化學習學習了一種基于偽度量（pseudometric）的動作表征。然而，MERLION 的策略執行需要在每個時間步對整個枚舉動作集進行最近鄰搜索，這在我們所考慮的組合場景中計算上是不可行的。此外，其架構將動作視為原子實體，并未對其底層的組合結構進行建模。相比之下，SPIN 專為這種組合設定設計，其結構化策略逐維度生成聯合動作，而非枚舉完整的組合動作集。

強化學習中的自監督預訓練。 強化學習中的自監督預訓練已呈現多種形式，包括作為表征塑造（representation shaping）的輔助目標（Jaderberg et al., 2016; Shelhamer et al., 2016）、對比與預測編碼器（Laskin et al., 2020; Schwarzer et al., 2021; Stooke et al., 2021; Liu & Abbeel, 2021b;a）以及世界模型建模（Ha & Schmidhuber, 2018）。其他研究探索了掩碼決策建模或軌跡建模（Cai et al., 2023; Liu et al., 2022; Wu et al., 2023; Sun et al., 2023）。大規模行為預訓練已催生出通用策略（generalist policies）與視覺-語言-動作模型（VLA models）（Brohan et al., 2022; Zitkovich et al., 2023; O’Neill et al., 2024; Kim et al., 2024; Team et al., 2024; Tirinzoni et al., 2025），并配套了預訓練后的快速適應方法（Sikchi et al., 2025）。這些方法大多以狀態或軌跡為中心，且通常預設了在線交互或多任務微調的場景。相比之下，SPIN 預訓練了一個能夠捕捉動作組合規律的動作結構模型（ASM），從而在無需任何在線交互的情況下，為組合動作空間中的策略學習提供結構化初始化。

3 預備知識

4 結構化策略初始化 (SPIN)

結構化策略初始化（Structured Policy INitialization, SPIN）是一個針對結構化動作空間中離線強化學習的兩階段框架，它顯式地將表示學習與控制解耦。在第一階段，通過自監督訓練一個動作結構模型（Action Structure Model, ASM），以學習一個表示函數；該函數以狀態 s 為條件，在子動作上誘導產生一個特征空間，其中結構上連貫的聯合動作集中在一個低維流形上。在第二階段，該表示被凍結，策略學習簡化為在誘導出的動作流形上訓練輕量級頭（heads），以用于下游的強化學習任務。

4.1 動作結構建模 (ASM)

ASM 的預訓練過程總結在算法 1 中。我們在附錄 C 中通過實證驗證了這一目標，展示了它優于強大的生成式和判別式替代方案。

4.2 基于凍結表示的策略學習

在第二階段，SPIN 在 ASM 提供的凍結表示上執行策略學習。策略網絡 π θ 僅更新輕量級組件，如查詢向量和輸出頭，而 ASM 保持固定。這種分離保留了學習到的動作結構，并保持策略優化的可處理性（tractability）。

5 實驗評估

為了隔離架構選擇的影響，所有方法均使用 IQL（Kostrikov et al., 2021）目標進行訓練。為了評估魯棒性，我們在附錄 D 中也報告了使用替代目標（包括 AWAC (Nair et al., 2020) 和 BCQ (Fujimoto et al., 2019)）的結果。為了驗證 SPIN 在運動（locomotion）之外的泛化能力，我們在 Maze（Beeson et al., 2024）上評估了其性能，結果見附錄 E。為了證明 SPIN 的有效性歸因于其以動作為中心（action-centric）的預訓練目標，而非僅僅源于預訓練本身，我們在附錄 F 中將其性能與一種以軌跡為中心（trajectory-centric）的預訓練方法進行了比較。在所有這些設置中，SPIN 在性能和效率上均一致優于基線方法。

所有實驗均使用 Python 3.9 和 PyTorch 2.6 在單塊 NVIDIA A40 GPU 上運行。報告的結果是五個隨機種子的平均值， ± ± 值表示跨種子的一個標準差。

5.1 漸近性能與訓練效率

表 1 報告了跨環境和數據集質量的最終性能與訓練效率（完整的學習曲線見附錄 A）。SPIN 取得了比所有基線方法 consistently 更高的回報，并且比所有基線方法用更少的掛鐘時間（wall-clock time）達到了目標性能。

SPIN 達到了最高的總體平均回報 594.1，超過了次優基線 SAINT 的 572.1。這種提升在整個基準測試套件中是系統性的，而非集中在個別環境中。這種優勢在異構的 medium-expert（中等 - 專家）和 random-medium-expert（隨機 - 中等 - 專家）數據集中最為顯著，它們代表了最現實且具有挑戰性的基準設置。在 random-medium-expert 數據集上，SPIN 實現了 499.2 的平均回報，比次優方法 SAINT（438.9）提升了超過 13%。

我們還測量了每種方法達到 F-IQL 漸近性能 95% 所需的掛鐘時間（以分鐘為單位報告）。F-IQL 是結構化動作空間中廣泛采用的最先進（state-of-the-art）基線（Tang et al., 2022; Beeson et al., 2024; Landers et al., 2024），在各環境中均提供了可處理性（tractability）和穩定的收斂性。使用 F-IQL 作為目標使得收斂到不同回報水平的方法之間能夠進行公平比較，避免了因在次優性能處提前終止而產生的誤導性優勢。我們采用 95% 的閾值而不是 100%，是因為有些方法永遠無法達到 F-IQL 的漸近性能。直接處理這些情況——無論是通過排除運行次數還是報告完整運行時間——都會使結果產生偏差，而 95% 的標準提供了一致且可比的度量。

每個環境的到達目標時間（time-to-target）完整結果報告在附錄 B 中。總計，SPIN 在 223.3 分鐘內達到目標性能，大約比 F-IQL 本身快 2.5 倍，比 SAINT 快 3.8 倍。這種加速在 medium-expert 數據集中尤為明顯，SPIN 僅需 62 分鐘的訓練時間，而所有其他方法則需要超過 250 分鐘。SPIN 的所有運行時間均包含 ASM 預訓練階段的全過程。

這些發現表明，在專用的預訓練階段顯式地對動作結構進行建模，使得表示層能夠捕捉連貫動作的流形。在策略學習期間凍結該表示保留了這種結構，使得輕量級頭（heads）能夠高效地適應下游任務。與 Factored（因子化）和 Autoregressive（自回歸）方法相比（它們要么丟棄跨維度依賴關系，要么對其施加強制性的剛性結構），SPIN 在保留靈活性的同時沒有犧牲可處理性。與試圖聯合學習動作結構和控制的 SAINT 不同，SPIN 的解耦設計實現了更高的漸近性能和更快的收斂速度。

5.2 對動作基數的魯棒性

結果總結在表 2 中。SPIN 在每個基數下都取得了最高的平均回報，且相對于基線的差距隨著動作空間的增大而增加。在三個區間時，SPIN 略微優于最強的基線 SAINT。在三十個區間時，SPIN 達到了 703.9 的平均回報，相比之下 SAINT 為 562.5，提升幅度超過 25%。AR-IQL 表現出不穩定的性能，從三個區間時的 526.5 下降到十個區間時的 457.4，而 F-IQL 則未顯示出從增加粒度中獲益，停留在 480 左右。

訓練效率遵循相同的趨勢。即使在最大的動作空間中，SPIN 始終需要更少的掛鐘時間（wall-clock time）來達到目標性能（完整運行時間結果見附錄 B）。這些結果表明，隨著組合復雜性的增長，將結構學習與控制分離開來日益有益，因為智能體可以在學習到的低維流形上行動，而端到端（end-to-end）的方法仍然受制于原始聯合空間的規模。

6 SPIN 有效性的底層機制

第 5 節的實驗表明，SPIN 在學習速度和最終性能上均優于現有方法。我們現在考察這些提升背后的機制。

6.1 表示質量對策略性能的影響

為了評估 ASM 預訓練的貢獻，我們在 medium-expert（中等 - 專家）數據集上將 ASM 表示訓練了 10–100 個 epoch。隨后，將每個表示函數凍結，并用于初始化一個新的策略，該策略隨后在控制任務上訓練至收斂。

圖 2 顯示，下游回報（return）通常隨著更多的 ASM 預訓練而提升，其中前 20 個 epoch 的增益最為陡峭。在 20 個 epoch 之后，策略在所有任務上均超過了完全收斂的 F-IQL 參考值。由未訓練的 ASM（Epoch 0）初始化的策略表現不佳。這些結果表明，最終策略性能在很大程度上取決于預訓練動作表示的質量；一旦學習到了連貫的表示，控制優化就會變得 substantially 更容易。

6.2 量化表示質量

圖 2 中隨機初始化（epoch 0）與預訓練智能體之間的巨大差距，可能是由于預訓練僅提供了方便的初始化而未編碼結構，也可能是由于預訓練學習到了能夠賦能下游性能的表示。我們通過測試 ASM 表示是否使用線性探針（linear probe）捕捉聯合動作依賴關系來直接評估這一點，線性探針是自監督表示的標準診斷工具（Chen et al., 2020; He et al., 2020）。

在本實驗中，ASM 表示被凍結——無論是預訓練了 100 個 epoch 還是隨機初始化——并在其嵌入上訓練一個輕量級線性分類器，以根據狀態預測數據集動作。為此探針（probe）學習了新的動作查詢和線性頭。分析是在 dog-trot 環境中進行的，該環境擁有 38 個子動作維度，被離散化為 30 個區間（bins），從而產生了 DM Control 套件中最大且最具挑戰性的組合動作空間。

6.3 分離學習到的表示的貢獻

訓練后，學生網絡被凍結，并作為下游策略的輕量級、無注意力機制的特征提取器發揮作用。表 3 報告了該實驗的結果。

SPIN-Distill 與完整 SPIN 模型的漸近性能相差無幾，并且顯著優于所有其他基線方法，同時速度比 SAINT 快近 8 倍。這些結果提供了強有力的證據，表明 SPIN 的性能提升歸因于預訓練表示本身的質量，而非策略網絡的具體架構。

6.4 涌現的快速適應

在確立了預訓練和表示質量的重要性之后，我們接下來考察學習動態。表 4 報告了在 10,000 個梯度步之后達到的 F-IQL 漸近性能的百分比，這僅相當于總訓練預算的 1%。在幾乎所有環境中，SPIN 學習到的策略都能達到至少 90% 的目標性能，而基線方法的提升則緩慢得多。這種效應在異構數據集上最為顯著。在使用 medium-expert 數據集的人形機器人（humanoid）任務中，SPIN 達到了目標性能的 93.4%，而次優方法 SAINT 僅達到了 9.3%。在 random-medium-expert 數據集上，在此期間，SPIN 在 cheetah 和 humanoid 任務中均超過了 F-IQL 漸近性能的 100%。

這種快速學習也闡明了 SPIN 的掛鐘時間效率（表 1）。下游 RL 階段的計算開銷主要由 Actor-Critic 循環占據，該循環需要在每個梯度步對 Actor、Critic 和目標網絡進行重復評估，以及進行貝爾曼備份（Bellman backups）。相比之下，ASM 預訓練階段是一個應用于掩碼子動作的、穩定的、單次遍歷（single-pass）的監督目標。因此，其相對成本極低：在 medium-expert 數據集上，預訓練僅占 cheetah 總掛鐘時間的 5.6%，finger 的 1.4%，以及 humanoid 和 quadruped 的 1.6%。

綜上所述，這些結果表明 ASM 提供了一個強有力的結構先驗，極大地簡化了下游學習。端到端基線方法必須聯合發現動作結構和控制，導致初始進展緩慢，而 SPIN 則利用連貫的表示開始策略學習，從而實現了高效的早期適應并減少了整體訓練時間。

7 討論與結論

在離散組合動作空間中進行強化學習，需要在指數級數量的組合動作中進行搜索，同時確保所選子動作構成連貫的動作組合。一些方法通過忽略動作結構來簡化策略學習（Tang et al., 2022; Beeson et al., 2024），但代價是丟棄了關鍵的子動作依賴關系。其他方法嘗試同時捕捉結構并求解控制問題（Zhang et al., 2018; Landers et al., 2024; 2025），但通常計算開銷極大且不穩定。相比之下，SPIN 采用兩階段過程將表示學習與策略學習分離。在第一階段，動作結構模型（ASM）學習一個表示函數，該函數以狀態 s s 為條件，在子動作上誘導一個特征空間，其中結構連貫的聯合動作位于一個低維流形上。隨后，該表示被凍結并在第二階段重用，此時控制問題簡化為在預訓練 ASM 之上訓練輕量級策略頭。

在數據集大小和質量、動作維度以及動作基數各異的基準測試中，SPIN 相比當前最優方法將平均回報最高提升了 39%，并將達到強基線性能所需的時間最多縮短了 12.8 倍。這些收益在具有挑戰性且更貼近現實的 medium-expert（中等-專家）和 random-medium-expert（隨機-中等-專家）數據集中最為顯著。

針對性分析闡明了 SPIN 的有效性。最終性能隨著學習到的表示質量的提升而提高，證實了控制問題的瓶頸在于結構發現。一旦該結構可用，策略便能快速學習，在極小的訓練比例內即可達到其最終回報的大部分。線性探針進一步表明，學習到的表示在生成完全協調的動作方面比隨機基線有效 45 倍，為下游智能體的成功提供了直接且定量的解釋。

盡管 SPIN 展現了強大的性能，但仍存在若干未來工作方向。將 SPIN 擴展至 CQL 等值正則化方法是一個有前景的方向。一個自然的下一步是開發混合目標，將 SPIN 的“表示優先”設計與溫和的保守正則化相結合——例如，將懲罰限制在 ASM 提議的候選聯合動作上，或應用于子動作級別，從而避免在完整組合空間上進行難以處理的全局操作。將 SPIN 適配于具有除排列等變性以外結構假設的動作空間（例如有序或基于序列的子動作）是另一個未來方向。最后，與所有離線方法一樣，SPIN 的泛化能力最終取決于數據集的覆蓋范圍，在稀疏或有偏數據下提高魯棒性仍然是一個重要的開放挑戰。

SPIN 為結構化動作空間中的控制引入了一種“表示優先”的視角。通過首先學習合理動作的流形，隨后重用表示函數進行下游決策，它將復雜的組合問題簡化為可處理的策略學習任務。這種解耦為高維、結構化領域的強化學習提供了一個原則性框架。

原文鏈接：https://arxiv.org/pdf/2601.04441