Technical Summary: Adaptive Chunk Size Actor-Critic with Causal Transformer Q-Network (ACSAC)

Problem Statement

Long-horizon, sparse-reward tasks present a fundamental challenge for reinforcement learning (RL). Standard single-step temporal difference (TD) learning suffers from bootstrapping error accumulation, where small errors compound across successive Bellman updates. While multi-step returns can accelerate value propagation, they often introduce off-policy bias because intermediate rewards are collected under a behavior policy rather than the current policy.

Action chunking has emerged as a solution, training critics and policies on action sequences rather than single actions. This allows for multi-step value backups without off-policy bias and promotes temporally coherent exploration. However, existing action-chunking methods rely on a fixed chunk size manually tuned per task. This rigidity forces a static trade-off between reactivity and temporal consistency:

• Large fixed chunks reduce responsiveness to new observations, potentially causing agents to "over-commit" to a trajectory in sensitive states (e.g., turns), driving them off the optimal path.
• Small fixed chunks produce incoherent motions and fail to leverage the benefits of long-horizon planning.

The optimal balance between reactivity and consistency is inherently state-dependent, yet current methods cannot adapt the execution horizon dynamically.

Methodology: ACSAC

The authors propose Adaptive Chunk Size Actor-Critic (ACSAC), a framework that adaptively selects the chunk size at each replanning state to maximize the expected return.

1. Causal Transformer Critic

ACSAC employs a causal Transformer critic to evaluate action chunks of varying lengths. Unlike standard MLP critics that require fixed-length inputs (leading to causal leakage when reusing weights on sub-prefixes), the Transformer ingests the state and a full-length candidate chunk $(s_t, a_{t:t+H})$ as a token sequence.

• Prefix-Conditioned Outputs: The network outputs $H$ values simultaneously: $\{Q_\phi(s_t, a_{t:t+h})\}_{h \in [H]}$, where each value corresponds to the expected return of executing exactly the first $h$ actions of the chunk.
• Causal Masking: A causal attention mask ensures the $h$-th output depends only on the prefix $a_{t:t+h}$, eliminating dependence on future actions in the chunk.
• Cross-Horizon Comparability: Because all $H$ values share a backbone and are trained against per-horizon targets on a common discounted-return scale, the Q-values for different prefix lengths are directly comparable.

2. Multi-Step TD Objective

The critic is trained using a multi-step TD loss that averages squared errors across all prefix lengths $h \in [H]$.

• Target Definition: For each prefix length $h$, the target is the sum of $h$ on-policy rewards from the dataset plus a bootstrap value from the target critic at the state $s_{t+h}$.
• Bootstrap Mechanism: The bootstrap value $Q_{\bar{\phi}}(s_{t+h}, \pi^*(s_{t+h}))$ is computed by sampling $N$ candidate chunks at the next state and selecting the prefix that maximizes the Q-value over all $N \times H$ candidates. This generalizes the EMaQ operator to a joint max over candidate indices and prefix lengths.
• Theoretical Guarantee: The authors prove that this Bellman operator is a $\gamma$-contraction with a unique fixed point corresponding to the action-value function of the adaptive policy.

3. Adaptive Policy Extraction

The actor utilizes a Flow Matching behavior cloning (BC) policy to generate expressive, multi-modal action chunks.

• Candidate Generation: At each replanning state $s_t$, the policy samples $N$ candidate chunks of maximum length $H$.
• Joint Rejection Sampling: Instead of fixing a chunk size, ACSAC performs a joint argmax over all $N$ candidates and all $H$ prefix lengths:
  $$(n^*, h^*) = \arg\max_{n \in [N], h \in [H]} Q_\phi(s_t, a^{(n)}_{t:t+h})$$
• Execution: The agent executes the optimal prefix $a^{(n^*)}_{t:t+h^*}$ and replans at the next boundary. This allows the agent to execute long chunks during stable phases (e.g., straight paths) and short chunks during sensitive phases (e.g., turns or grasping), balancing reactivity and coherence without manual tuning.

Key Contributions

1. Adaptive Chunk Selection: Introduction of ACSAC, the first action-chunking actor-critic that adaptively selects chunk sizes per state using a causal Transformer critic and joint rejection sampling.
2. Theoretical Foundations: Proof that the prefix-conditioned Q-values are cross-horizon comparable and that the ACSAC Bellman operator is a contraction whose unique fixed point is the value function of the adaptive policy.
3. Empirical Performance: Demonstration that ACSAC achieves state-of-the-art performance on long-horizon, sparse-reward manipulation tasks (OGBench) in both offline and offline-to-online settings, outperforming single-step, multi-step, and fixed-chunk baselines.

Experimental Results

The method was evaluated on the OGBench suite, comprising 25 tasks across five domains (scene-sparse, puzzle-3x3-sparse, cube-double, cube-triple, cube-quadruple).

• Overall Performance: ACSAC consistently outperformed all baselines (including IQL, ReBRAC, FQL, BFN, and fixed-chunk methods like QC and DEAS) in both offline pre-training and online fine-tuning phases.
• Long-Horizon Tasks: The improvements were most significant in domains requiring a mix of coarse transport and precise manipulation (e.g., cube-triple and scene-sparse), where adaptive chunking allowed the agent to switch execution horizons based on task phases.
• Qualitative Analysis: Visualizations of the executed chunk sizes showed that ACSAC naturally aligned with task semantics: selecting longer chunks for coarse approach phases and progressively shorter chunks for precise placement phases.
• Calibration: Experiments confirmed that the prefix-conditioned Q-values were well-calibrated against realized returns and comparable across different horizon lengths, validating the joint argmax strategy.
• Ablation Studies:
  • Chunk Size ($H$): Single-step execution ($H=1$) failed on sparse-reward tasks, while overly long maximum chunks ($H=7$) degraded performance in some domains.
  • Sampling Size ($N$): Increasing the rejection sampling size beyond $N=4$ yielded diminishing returns.
  • Architecture: Replacing the MLP critic in baselines with the causal Transformer alone improved performance but did not match ACSAC, confirming that the adaptive selection mechanism and multi-horizon objective are critical.

Significance and Claims

The paper claims that ACSAC addresses the fundamental limitation of fixed chunk sizes in action-chunking RL. By leveraging a causal Transformer to learn cross-horizon comparable Q-values, ACSAC enables state-dependent execution horizons. This allows agents to retain the fast value backups and temporal consistency of long chunks while maintaining the reactivity required for precise control in sensitive states.

The authors position ACSAC as a step toward more robust RL for real-world applications, noting that future work could integrate these state-dependent chunk sizes with large vision-language-action (VLA) models, building on recent trends where chunked value learning scales to complex multimodal settings. The method is presented as a unified approach for both offline and offline-to-online learning, requiring no specialized techniques beyond adding online transitions to the dataset.