HWC
← Note

No date

Trajectory Transformer | Offline Reinforcement Learning as One Big Sequence Modeling Problem

PaperRLOffline RLWorld Model

快速解釋

這篇論文的核心主張是：offline RL 不一定要拆成 policy learning、value learning、或 single-step dynamics modeling，也可以直接把整條軌跡當成一個大型序列建模問題。作者把 state、action、reward，甚至 reward-to-go 全都離散化成 token，然後用 GPT 風格的 Transformer 去學習整個 trajectory 的 joint distribution。推論時不直接抽樣單步 action，而是像語言模型解碼一樣，用 beam search 搜尋未來幾步最有希望的軌跡，再執行第一個 action 並重複規劃。

這個觀點的關鍵價值在於：同一個模型同時扮演 dynamics model、behavior prior 與規劃器的基礎機率模型。因為模型是聯合建模 state 與 action，所以在規劃時天生偏向資料分佈內的行為，不需要像許多 offline RL 方法那樣額外設計 pessimism、Q regularization 或 explicit behavior constraint。這篇論文後來也因為這個角度，被廣泛稱為 Trajectory Transformer。

問題設定

作者考慮的是標準的 offline reinforcement learning。給定一個只能讀取、不能再與環境互動的資料集：

D={τ(k)}k=1K\mathcal{D} = \{\tau^{(k)}\}_{k=1}^K

其中每條軌跡由 state、action、reward 組成：

τ=s1,a1,r1,s2,a2,r2,…,sT,aT,rT\tau = s_1, a_1, r_1, s_2, a_2, r_2, \ldots, s_T, a_T, r_T

目標是在不做額外 online exploration 的情況下，從資料中學到一個能產生高回報行為的控制器。

傳統 offline RL 常見做法有兩類。第一類是 model-free，透過 Bellman backup 學 Q-function，再從 Q-function 萃取 policy；第二類是 model-based，先學單步 dynamics，再做 trajectory optimization。這兩類方法都高度依賴 Markov factorization，也通常需要額外處理 distribution shift，例如限制 policy 不要偏離 dataset support，或在 value function 中加入 pessimism。

這篇論文提出的問題是：如果我們完全不把 RL 問題分解成 value estimation、policy extraction、或 single-step prediction，而是直接學習「整條軌跡的分佈」，那麼 sequence modeling 本身能不能就是一種 RL 演算法？

為了讓 Transformer 處理連續控制資料，作者把每個 state 維度與 action 維度獨立離散化。若 state 維度是 N、action 維度是 M，則每個 timestep 會展開成一串 token：

τ=…,st1,st2,…,stN,at1,at2,…,atM,rt,…\tau = \ldots, s_t^1, s_t^2, \ldots, s_t^N, a_t^1, a_t^2, \ldots, a_t^M, r_t, \ldots

因此原本的控制問題被轉寫成標準的自回歸序列建模問題。

核心想法

作者的方法可以拆成四個互相配合的設計。

  1. 把 trajectory 當成語言序列來做 joint modeling

與其只學單步轉移 p(st+1midst,at)p(s_{t+1} mid s_t, a_t)，作者直接學整個軌跡的自回歸分佈。這意味著模型不只是在預測下一個 state，也在同時建模哪些 action、哪些 reward pattern 與哪些未來狀態會共同出現。這讓模型具備比傳統 Gaussian single-step model 更強的表達能力，也更能捕捉長期相依關係。

  1. 規劃不是 value backup，而是 sequence decoding

訓練完模型後，控制問題被改寫成 decoding 問題。給定目前觀測到的 state 與歷史，模型會往未來展開多個候選軌跡；beam search 保留最有希望的候選，最後從最佳軌跡取出第一個 action 執行。也就是說，Trajectory Transformer 並不是直接輸出一個 stationary policy，而是在每個時間點做一次基於模型的 receding-horizon planning。

  1. offline RL 版本用 reward-to-go 避免 beam search 太短視

如果 beam search 只按 likelihood 排序，它會偏向模仿資料中高機率的行為；如果只按 immediate reward 排序，又容易變得短視。因此作者在每個 transition 後面加入 reward-to-go，讓模型除了預測當下 reward，也預測從該時刻開始的未來折扣回報。規劃時，beam search 可以利用「累積 reward + reward-to-go」來排序候選軌跡，等於把 Monte Carlo value estimate 當成 heuristic。

  1. in-distribution constraint 不是顯式加上去，而是隱含在 joint model 裡

這篇方法最有意思的地方，是它不顯式估計 behavior policy，也不需要像 CQL、BRAC 那樣額外寫一個保守化目標。因為 action 不是獨立優化變數，而是和 state 一起由 joint trajectory model 生成，模型在搜尋時自然偏向資料集中曾經共同出現過的 state-action pattern。作者認為這能減少 model exploitation 與 out-of-distribution action 的問題。

另外，論文還展示了同一套模型可直接延伸到 imitation learning 與 goal-conditioned RL。對 imitation learning，只要做 likelihood-maximizing beam search；對 goal-conditioned RL，只要把未來 goal state 當成條件輸入即可。也就是說，差別主要在 decoding objective，而不是模型本身。

關鍵公式

Trajectory Transformer 的基本訓練單位是一條 trajectory：

τ=s1,a1,r1,…,sT,aT,rT\tau = s_1, a_1, r_1, \ldots, s_T, a_T, r_T

對連續控制任務，作者先把 state 與 action 的每個維度各自離散化，因此在 token 層級上，每個 timestep 會變成：

(st1,…,stN,at1,…,atM,rt)(s_t^1, \ldots, s_t^N, a_t^1, \ldots, a_t^M, r_t)

訓練時使用標準的 teacher forcing，自回歸地最大化整條序列的對數似然。論文中的目標函數寫成：

L(τ)=∑t=1T[∑i=1Nlog⁡Pθ ⁣(sti∣st<i,τ<t)+∑j=1Mlog⁡Pθ ⁣(atj∣at<j,st,τ<t)+log⁡Pθ ⁣(rt∣at,st,τ<t)]L(\tau) = \sum_{t=1}^{T} \Biggl[ \sum_{i=1}^{N} \log P_{\theta}\!\left(s_t^i \mid s_t^{<i}, \tau_{<t}\right) + \sum_{j=1}^{M} \log P_{\theta}\!\left(a_t^j \mid a_t^{<j}, s_t, \tau_{<t}\right) + \log P_{\theta}\!\left(r_t \mid a_t, s_t, \tau_{<t}\right) \Biggr]

這個式子的意思是：模型先依序預測 state 各維度，再預測 action 各維度，最後預測 reward；所有 prediction 都建立在先前 token 與歷史 timestep 的條件上。

為了讓 offline RL 的規劃不只看 immediate reward，作者額外把 reward-to-go 加進 trajectory 中：

Rt=∑t′=tTγ t′−trt′R_t = \sum_{t' = t}^{T} \gamma^{\,t' - t} r_{t'}

因此在離線規劃時，模型實際上學到的是包含 state、action、reward、reward-to-go 的擴充 transition。這讓 beam search 在評估候選序列時，不只考慮目前累積到的 reward，也能利用對未來回報的估計作為啟發式。

beam search 的一般形式是：先擴展候選集合

Ct={yt−1∘y∣yt−1∈Yt−1,  y∈V}C_t = \{ y_{t-1} \circ y \mid y_{t-1} \in Y_{t-1},\; y \in \mathcal{V} \}

再從中保留 beam width 為 B 的最佳候選：

Yt=arg⁡max⁡Y⊆Ct,  ∣Y∣=Blog⁡Pθ(Y∣x)Y_t = \arg\max_{Y \subseteq C_t,\; |Y| = B} \log P_{\theta}(Y \mid x)

在 imitation learning 中，排序依據主要是序列 likelihood；在 offline RL 中，作者把 beam search 改成偏好高 reward 的候選，並用 reward-to-go 補足長期價值資訊。若把這個規劃分數寫成簡化形式，可以理解成：

Score=∑k=tt+H−1rk+Rt+H\text{Score} = \sum_{k=t}^{t+H-1} r_k + R_{t+H}

其中第一項是 rollout 中已累積的 reward，第二項是模型預測的未來回報。這不是訓練目標，而是規劃時用來排序候選軌跡的啟發式。

論文也給出 goal-conditioned decoding 的條件分佈形式。若把最終目標 state 設為條件，則可以建模：

Pθ ⁣(sti∣st<i,τ<t,sT)P_{\theta}\!\left(s_t^i \mid s_t^{<i}, \tau_{<t}, s_T\right)

這表示在已知未來目標的條件下，模型可以反推出一條最有可能到達該目標的軌跡。

模型結構

這篇論文的「模型結構」可以分成資料表示、Transformer 骨幹、以及規劃機制三層來看。

  1. 資料表示

作者對每個 state 與 action 維度做獨立離散化，並比較兩種 tokenization 方式：uniform discretization 與 quantile discretization。uniform 保留原始空間中的幾何距離感；quantile 則確保每個 token 都有足夠資料覆蓋。實驗中，quantile 在某些高動態範圍任務上更穩定。

  1. Transformer 骨幹

模型本體是 GPT 風格的 decoder-only Transformer。根據 appendix，作者使用的是 4 層 Transformer、4 個 self-attention heads、token embedding dimension 為 128，block 後 dropout 為 0.1。輸入 vocabulary 的總大小對應 state、action、reward、reward-to-go 各類 token；訓練時用 Adam，learning rate 線性 warmup 到 2.5times10−42.5 times 10^{-4}，batch size 為 256。

  1. 自回歸預測順序

在單一 timestep 內，模型不是一次輸出整個向量，而是像語言模型一樣逐 token 預測：先輸出 state 各維度，再輸出 action 各維度，接著是 reward 與 reward-to-go。這個設計讓模型能捕捉高維連續控制資料中的複雜 joint dependency，而不是預設高斯或對角協方差分佈。

  1. 規劃與控制迴圈

推論時，給定當前 state 與過去 context，beam search 往前展開固定 planning horizon 的多條候選軌跡。每個 timestep 最後只執行最佳候選的第一個 action，然後重新觀測、重新規劃，所以本質上是 model-predictive control。論文 appendix 中給出的典型設定包含：beam width 128、planning horizon 15、vocabulary size 100、context size 5。

  1. 稀疏回饋版本：TT + Q

在稀疏回饋的 AntMaze 任務中，單靠 Monte Carlo reward-to-go 不夠提供有效搜尋訊號，因此作者把 Trajectory Transformer 與 IQL 的 Q-function 結合，改用 Q-function 作為 beam search 的 heuristic。這個版本保留 sequence planning 的長期組合能力，同時利用 dynamic programming 補強稀疏回饋下的價值估計。

總結來說，Trajectory Transformer 並不是把 Transformer 當成一個 policy network，而是把它當成「軌跡機率模型」，再把控制問題外移到 decoding 與搜尋。這也是這篇論文最重要的觀念轉換。