v3/paper_source/04_method_short.tex

\begin{figure*}[t]
    \centering
    \vspace{-2mm}
    \label{fig:method_overview}
    
    \includegraphics[width=0.95\linewidth]{figures/method_overview.pdf}
    % \includegraphics[width=0.6\linewidth]{./doc/benchmark-demo-v7.pdf}
    \caption{\textbf{The training process of SuP.} Our method is trained purely on offline datasets through three phases: (1) recurrent world model learning; (2) data synthesis; and (3) scheduler optimization via IQL. Through this pipeline, we optimize the scheduler to achieve the maximum possible speedup while preserving comparable performance.}
    \vspace{-3mm}
\end{figure*}

To solve the CMDP for adaptive manipulation acceleration, we propose Speedup Patch (SuP), a dual-stage framework. Our objective is to learn an optimal scheduler policy $\pi_\phi$ that maximizes execution speed while maintaining task fidelity from an offline demonstration dataset $\mathcal{D}$. We first train a Recurrent World Model (RWM) on $\mathcal{D}$ (Sec. \ref{sec:world_model}), allowing us to evaluate potential state deviations and constraint violations. Using the RWM as a data evaluator and generator, we synthesize a CMDP dataset. We then optimize the final scheduler via Implicit Q-Learning (Sec. \ref{sec:scheduler}).

\subsection{Recurrent World Model}
\label{sec:world_model}
To estimate the state deviation $\mathcal{E}$ and the violation signal $h$ in an offline setting, we develop a Recurrent World Model (RWM), denoted as $\mathcal{M}_\theta$. To ensure the framework remains lightweight, $\mathcal{M}_\theta$ is designed to predict the robot's state $o$ directly, thereby bypassing the high-dimensional reconstruction of visual observations $I$.

Our RWM architecture is built upon ADM \cite{lin2024any}, which is specifically capable of handling variable-length action sequences while mitigating the compounding errors typical in multi-step rollouts. A pivotal feature of this design is that predicted states $\hat{o}$ are never fed back as inputs for subsequent time steps. Instead, the model evolves its hidden state exclusively through the action sequence. 

The RWM is optimized via variable-length sequence supervision. During training, we sample action sequences $A_t$ of variable length $L \in [1, L_{max}]$ from the offline dataset $\mathcal{D}$, where the range of $L$ is chosen to encompass the temporal scales generated by different downsampling factors $k$. The parameters $\theta$ are updated to minimize the multi-step Mean Squared Error:
\begin{equation}\label{eq:wm_loss}\mathcal{L}(\theta) = \mathbb{E}_{(o_t, A_t, o_{t+1:t+L}) \sim \mathcal{D}} \left[ \sum_{i=1}^{L} \|\hat{o}_{t+i} - o_{t+i}\|^2_2 \right].\end{equation}

\subsection{Scheduler Optimization}
\label{sec:scheduler}
With the world model $\mathcal{M}_\theta$, we now describe the learning process for the scheduler policy $\pi_{\phi}$. To solve the CMDP, we first transform the constrained problem into an unconstrained MDP via a penalty-augmented reward:
\begin{equation}
\label{eq:reward_new}
r'(s,k) = \begin{cases} k, & h_\mathcal{E}(s,k)=0 \\ -\Omega, & h_\mathcal{E}(s,k)=1 \end{cases}
\end{equation}
where $\Omega$ is a sufficiently large penalty to guarantee zero-violation safety (Prop. \ref{prop:penalty}). To facilitate offline training, we construct a synthetic dataset $\mathcal{D}'$ by re-labeling the original demonstrations. Specifically, for each transition in $\mathcal{D}$, we use $\mathcal{M}_\theta$ to evaluate the violation signal $h$ across all potential downsampling factors $k$, generating a rich set of counterfactual transitions $(o, A^k, r', o')$.

\begin{proposition}
\label{prop:penalty}
Let $K_{\max}$ be the maximum possible speedup rate and $\gamma \in [0, 1)$ be the discount factor. If the penalty $\Omega$ satisfies the condition:
\begin{equation}
\Omega > \frac{\gamma K_{\max}}{1-\gamma},
\end{equation}
then the optimal policy $\pi^*$ maximizing the cumulative reward of $r'(s, k)$ satisfies the constraint $h(s, \pi^*(s)) = 0$ for all reachable states $s$. See App. \ref{app:proof2} for the proof.
\end{proposition}

In a standard RL setting, computing temporal difference (TD) errors requires the next action $A'$, which is unavailable since we do not know how expert will behave with $s'$. We resolve this by transforming the policy input: instead of directly processing the raw tuple $(s, k)$, the scheduler maps the action chunk $A$ and skip-length $k$ into a downsampled representation $A^k$. By conditioning the policy on $(o, A^k)$, we can treat the resulting $s'$ directly as the subsequent state in a Markovian transition, effectively bypassing the need for future action chunks during value estimation.

Finally, we employ Implicit Q-Learning (IQL) \cite{kostrikovoffline} to optimize the scheduler on $\mathcal{D}'$. The value function $V_\psi$ and Q-function $Q_\phi$ are learned via expectile regression:\begin{equation}
\label{eq:iql}
\begin{aligned}
    L_Q(\phi)&=\mathbb{E}_{(o,A^k,r',o')\sim D'}[(r'+\gamma V_\psi (o')-Q_\phi(o,A^k))^2],\\
    L_V(\psi)&=\mathbb{E}_{(o,A^k)\sim D'}[L^\alpha_2( V_\psi (o)-Q_\phi(o,A^k))],
\end{aligned}
\end{equation}where $L^\alpha_2(x)=|\alpha-\mathbb{I}(x<0)|x^2$ is the expectile loss. During inference, the optimal skip-length is determined by $\pi_\phi(o,A)=\arg\max_k Q_\phi(o,A^k)$.

\begin{algorithm}[t]
   \caption{Training Procedure of SuP}
   \label{alg:sup_training}
\begin{algorithmic}[1]
   \STATE {\bfseries Input:} Offline demonstration dataset $\mathcal{D}$, minimum and maximum downsampling rate $k_{\min},k_{\max}$ penalty $\Omega$, deviation threshold $\epsilon$.
   \STATE {\bfseries Output:} Scheduler policy $\pi_{\phi}$.
   
   \STATE \COMMENT{\textbf{Phase 1: Recurrent World Model Learning}}
   \STATE Initialize world model $\mathcal{M}_\theta$.
   \WHILE{not converged}
      \STATE Sample batch of data $(o_t, A_t, o_{t+1:t+L})$ from $\mathcal{D}$.
      \STATE Update $\mathcal{M}_\theta$ with Eq. \ref{eq:wm_loss}.
   \ENDWHILE

   \STATE \COMMENT{\textbf{Phase 2: Data Synthesis}}
   \STATE Initialize synthetic dataset $\mathcal{D}' \leftarrow \emptyset$.
   \FOR{each transition $(o_t, A_t)$ in $\mathcal{D}$}
      \FOR{$k = k_{\min}$ {\bfseries to} $k_{\max}$}
         \STATE Construct downsampled action chunk $A^k_t$.
         \STATE Predict next states $\hat{o}'_{t+1:t+L} \leftarrow \mathcal{M}_\theta(o, A^k)$.
         \STATE Estimate deviation $\mathcal{E}$ with Eq. \ref{eq:state_deviation} and violation signal $h_\mathcal{E} \leftarrow \mathbb{I}(\mathcal{E} > \epsilon)$.
         \STATE Compute reward $r'_t$ with Eq. \ref{eq:reward_new}.
         \STATE Store transition $(o_t, A^k_t, r'_t, \hat{o}'_{t+L})$ into $\mathcal{D}'$.
      \ENDFOR
   \ENDFOR

   \STATE \COMMENT{\textbf{Phase 3: Scheduler Optimization via IQL}}
   \STATE Initialize IQL networks $V_\psi, Q_\phi$.
   \WHILE{not converged}
      \STATE Sample batch $(o, A^k, r', o')$ from $\mathcal{D}'$.
      \STATE Update $V_\psi,Q_\phi$ with Eq. \ref{eq:iql}.

   \ENDWHILE
   \STATE \textbf{Return} Scheduler $\pi_\phi(o, A) = \arg\max_k Q_\phi(o, A^k)$.
\end{algorithmic}
\end{algorithm}
\vspace{-5mm}