Chunking the Critic: A Transformer-based Soft Actor-Critic with N-Step Returns

2025/03/05 by Tian, Dong, Celik, Onur, Neumann, Gerhard · 5 citations
#FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2503.03660

Abstract

We introduce a sequence-conditioned critic for Soft Actor--Critic (SAC) that models trajectory context with a lightweight Transformer and trains on aggregated N-step targets. Unlike prior approaches that (i) score state--action pairs in isolation or (ii) rely on actor-side action chunking to handle long horizons, our method strengthens the critic itself by conditioning on short trajectory segments and integrating multi-step returns -- without importance sampling (IS). The resulting sequence-aware value estimates capture the critical temporal structure for extended-horizon and sparse-reward problems. On local-motion benchmarks, we further show that freezing critic parameters for several steps makes our update compatible with CrossQ's core idea, enabling stable training without a target network. Despite its simplicity -- a 2-layer Transformer with 128-256 hidden units and a maximum update-to-data ratio (UTD) of 1 -- the approach consistently outperforms standard SAC and strong off-policy baselines, with particularly large gains on long-trajectory control. These results highlight the value of sequence modeling and N-step bootstrapping on the critic side for long-horizon reinforcement learning.

Cited by

• Multistep Quasimetric Learning for Scalable Goal-conditioned Reinforcement Learning
• DEAS: DEtached value learning with Action Sequence for Scalable Offline RL
• Actor-Critic for Continuous Action Chunks: A Reinforcement Learning Framework for Long-Horizon Robotic Manipulation with Sparse Reward
• Sequence Aware SAC Control for Engine Fuel Consumption Optimization in Electrified Powertrain
• CO-RFT: Efficient Fine-Tuning of Vision-Language-Action Models through Chunked Offline Reinforcement Learning

Related

• Actor-Critic for Continuous Action Chunks: A Reinforcement Learning Framework for Long-Horizon Robotic Manipulation with Sparse Reward
• SACn: Soft Actor-Critic with n-step Returns
• Improving Exploration in Soft-Actor-Critic with Normalizing Flows Policies
• Improving Actor-Critic Training with Steerable Action-Value Approximation Errors
• PAC-Bayesian Soft Actor-Critic Learning