Three Steps at a Time: Learning Representations from Action Sequences in Contrastive RL

Michal Korniak; Kamil Dybek; Benjamin Eysenbach; Marco Bagatella; Michał Bortkiewicz

Source-linked AI summary

Three Steps at a Time: Learning Representations from Action Sequences in Contrastive RL

Michal Korniak, Kamil Dybek, Benjamin Eysenbach, Marco Bagatella, Michał Bortkiewicz

arXiv:2608.30640v1Aug 31, 2026cs.LG

TL;DR

The paper asks whether CRL should model action sequences rather than single-step actions. It extends CRL with action-chunk conditioning and finds large gains across offline and online benchmarks, with chunk-based goal information improving critic representations. The benefit is conditional on chunk length because very long chunks can hinder policy extraction.

Problem
CRL traditionally models whether a goal will be visited after a single action, leaving the useful time scale for action modeling an open question.
Method
CRL + AC conditions the critic on an action chunk rather than a single action and modifies the actor to output a full chunk.
Results
+31.7% average improvement occurs on manipulation tasks and +93.1% on online locomotion and navigation tasks, with gains consistent across settings.
Takeaways & Limitations
Action chunks provide the CRL critic more information about goals than single actions, improving representations and algorithm effectiveness.
Takeaways & Limitations
Performance depends on selecting an appropriate chunk length; the authors generally find H=3 most beneficial, while very long chunks can hinder policy extraction.

Abstract

from arXiv · show

While self-supervised approaches to reinforcement learning have achieved strong results by learning representations of states and actions, a key open question is the time scale over which actions should be modeled. Departing from the standard formulation relying on single-step actions, we extend contrastive reinforcement learning (CRL), a prototypical self-supervised method, to operate over action chunks, and find that this results in large, pervasive gains across established offline and online benchmarks: +31.7% and +93.1% across 18 and 11 environments respectively. While action-chunking-driven gains are generally explained through the ability to model non-Markovian, temporally extended policies, and to propagate unbiased multi-step returns, interestingly, we find that these arguments only partially apply to CRL. Our empirical studies suggest that, in the context of CRL, an action chunk carries more information about the goal than a single action, measurably improving the critic's representations, and rendering the algorithm significantly more effective.

1 INTRODUCTION

The paper extends CRL by conditioning its critic on action sequences rather than single actions, yielding consistent gains across offline and online benchmarks. The authors attribute these gains partly to action chunks carrying more goal information for the critic.

CRL learns representations by predicting whether a goal will be visited after starting from a state and executing a single action.
The proposed modification conditions the CRL classification problem on action sequences to learn stronger critic representations.
Action chunking consistently improves performance across offline manipulation, noisy and exploratory datasets, and online locomotion and navigation tasks.
Action chunks give the CRL critic more information about the goal than single actions, improving representations and making the algorithm more effective, especially on noisy datasets.
+31.7% average improvement is reported on the OGBench manipulation suite, while +69.4% is reported on suboptimal OGBench datasets.
+93.1% average improvement is reported on online JaxGCRL locomotion and navigation tasks.

2 RELATED WORK

The related work situates the paper in goal-conditioned reinforcement learning, contrastive representation learning, and action chunking. Its focus is controlling action representations within CRL.

Goal-conditioned RL and action representations: Goal-conditioned reinforcement learning conditions policies and critics on goals to achieve varied tasks.
Contrastive Learning and Contrastive RL: Contrastive reinforcement learning casts goal-conditioned reinforcement learning as representation learning that distinguishes likely future states from unlikely ones.
Contrastive Learning and Contrastive RL: Prior CRL research emphasizes representation scalability and exploration, while this work studies gains from controlling action representations.
Action Chunking: Action chunking represents a policy output as H consecutive atomic actions instead of a single action, with open-loop and closed-loop execution variants.

3 BACKGROUND

CRL formulates goal-conditioned reinforcement learning as contrastive representation learning. It learns state-action and goal encoders whose compatibility score represents future-goal reachability and supports action ranking.

The environment is modeled as a reward-free Markov decision process with state and action spaces, transitions, an initial-state distribution, and discount factor.
The successor measure gives the probability of visiting a future state under a goal-conditioned policy and is equivalent to a goal-conditioned Q-function.
CRL trains state-action and goal encoders so their inner product recovers a log-probability ratio between positive future-state samples and negative dataset samples.
The contrastive objective lower bounds mutual information, and at its optimum exponentiated representation dot products equal Q-values up to an action-independent constant.
The learned representations can rank actions at each state, and the policy is trained to maximize their state-action–goal dot products.

4 ACTION-CHUNKED CONTRASTIVE RL

CRL + AC extends CRL by conditioning the critic on an H-action chunk and modifying the actor to output the same chunk. The replay and sampling pipeline remains largely unchanged, while inference can replan before executing the full chunk.

CRL + AC conditions the critic on a sequence of H future actions instead of a single action.
The positive goal is sampled geometrically from future trajectory states independently of H, so it need not be the state H steps ahead.
The actor is modified to output a full action chunk and uses the same policy-improvement objective as standard CRL.
The replay buffer remains unchanged; at sampling time, H consecutive stored actions are concatenated and a future goal is sampled geometrically.
At inference, the policy executes H_exec actions from each predicted chunk before replanning, with H_exec=H giving open-loop execution and H_exec=1 giving stepwise replanning.
Only the critic input and actor output change, and different policy-extraction procedures can be applied on top of the chunked critic.

5 EXPERIMENTS

Across offline and online benchmarks, action-chunked CRL consistently improves performance over standard CRL. The gains reflect both better critic representations from goal-relevant action information and additional action-chunking effects, while overly long chunks can make policy extraction the bottleneck.

5.1 RESULTS: +111.2% average time-at-goal on online locomotion and navigation tasks, with gains consistent across the evaluated task suite.CRL + AC with H ∈{3, 5} consistently outperforms CRL, and H=3 is best or near-best across most tasks.
5.1 RESULTS: +31.7% average performance on offline manipulation tasks, rising to +69.4% on noisy and exploratory datasets.Performance on play environments remains on par with standard CRL, without degrading any task group.
5.2 WHY DOES ACTION CHUNKING HELP?: Distillation shows that improved critic representations contribute to performance gains, with distilled single-step policies outperforming standard CRL by ∼30% to ∼100% depending on the dataset.The full-chunk policy improves further, attributed to non-Markovian modeling in this setting.
5.2 WHY DOES ACTION CHUNKING HELP?: Conditioning the critic on action chunks adds more information about the goal than conditioning on a single action, improving its representations.Validation accuracy gains over a state-only critic reach ∼9.5 percentage points at H=30, nearly double the ∼4.5-point gain at H=1.
5.2 WHY DOES ACTION CHUNKING HELP?: At large chunk lengths, critic quality and policy quality diverge: accuracy peaks at H=30 while success peaks at H=5 and collapses near zero by H=30.With single-step replanning, distilling long-chunk critics into an Hπ=5 critic recovers success to ∼50%, implicating policy extraction rather than critic quality.
5.3 ABLATION STUDIES: Action chunking remains advantageous across network depths, is comparable under open-loop execution and replanning, and becomes more valuable as dataset noise increases.On Ant Hardest Maze, a 2-layer CRL + AC network matches a 16-layer CRL network; noise amplification remains unexplained by the critic metrics.

6 CONCLUSIONS

The paper shows that short action chunks improve CRL across offline and online benchmarks by enriching critic representations with goal-relevant information. Very long chunks can improve critics while hindering policy extraction, and selecting the appropriate horizon remains necessary.

Mechanism: Action chunks give the CRL critic more information about the goal than single actions, improving representations and algorithm effectiveness.This mechanism is presented as specific to CRL and complementary to standard action-chunking explanations.
Horizon effects: Very long chunks improve critic quality but hinder policy extraction, while decoupling critic and policy horizons largely recovers performance.The paper uses separate horizons to retain representational benefits without requiring direct extraction from excessively long chunks.
Limitations and future work: The benefits remain conditional on choosing an appropriate chunk length, with H=3 generally found most beneficial in practice.Automatic horizon selection and alternative temporal enrichments remain open directions.

C.3 FIGURES

The figures summarize offline and online CRL + AC experiments across manipulation, noisy or exploratory, locomotion, and navigation environments. Reported aggregates use fixed evaluation protocols, confidence intervals, and comparisons between standard CRL and action-chunked CRL.

Aggregate coverage: CRL + AC uses H=3 and aggregates 18 manipulation environments, 12 noisy or explore environments, and 11 online locomotion or navigation environments.The visual OGBench environments are excluded, leaving state-based observations for the reported aggregates.
Offline results: 32% relative improvement is reported over manipulation environments, with gains increasing to 69.4% on noisy and explore variants.The values are RLiable aggregate relative improvements based on success rate comparisons.
Evaluation reporting: Table 1 reports success rate at the end of training, averaged over evaluation goals and 3 seeds with 95% confidence intervals.Group aggregates and intervals are computed with RLiable using stratified bootstrap over seeds and tasks.

D.1 OGBENCH EVALUATION PROTOCOL

The OGBench protocol evaluates policies repeatedly during training and reports task-aggregated success rates with bootstrapped confidence intervals. The experiments cover listed offline manipulation environments alongside listed JaxGCRL environments.

OGBench evaluation: OGBench policies are evaluated every 50,000 training steps across 20 checkpoints over 1,000,000 training steps.Each evaluation rolls out 20 episodes across 5 tasks per environment.
Reporting: OGBench results report mean task-aggregated success rates averaged over 3 seeds with 95% bootstrapped confidence intervals.All runs and seeds are included in the reported results.
Environments: The listed OGBench environments include cube manipulation and puzzle tasks in play and noisy variants.Examples include cube-single-play-v0, cube-single-noisy-v0, puzzle-3x3-play-v0, and puzzle-4x4-play-v0.
Environments: The listed JaxGCRL environments include Ant locomotion and multiple antmaze exploration tasks.The list includes Ant variants, antmaze-medium-explore-v0, antmaze-large-explore-v0, and antmaze-teleport-explore-v0.

E.1 ANALYSIS: DECOUPLING CRITIC AND POLICY HORIZONS (DQC)

DQC separates the representational quality of a chunked critic from the horizon used for policy extraction. Distillation shows that information from well-trained long-chunk critics can remain usable even when direct policy extraction collapses.

DQC method: DQC distills a long-chunk critic into a short-chunk critic by training an auxiliary encoder with expectile regression.The long-chunk critic remains trained by the contrastive objective, while only the auxiliary encoder receives the distillation loss.
DQC settings: DQC uses κ=0.9 and Hπ=5 for large-horizon recovery, versus κ=0.95 and Hπ=1 for single-action distillation.All other hyperparameters match the main results.
Single-step distillation: A distilled single-action policy from an H=5 critic outperforms standard CRL despite conditioning on and executing only one action.This isolates a representational benefit that does not require non-Markovian policy modeling or multi-step execution.
Large-horizon recovery: H=30 and H=50 critics remain near 0% when extracted directly, but distillation into Hπ=5 recovers approximately 50% success for both.The result indicates that the long-chunk critics contain usable information even when direct extraction fails.

E.2 VARYING NOISE ABLATION

Across noisy offline datasets, action chunking improves critic discrimination and performance, with larger gains when noise makes state–action discrimination harder. The benefit also persists on tasks where state information remains stable across noise levels.

Critic discrimination: Action chunking improves validation accuracy and lowers contrastive loss across noise levels, from state-only CRL through CRL + AC with H=3 and H=5.The progression reflects increasing action information in the critic input.
Noise dependence: On cube, noise sharply degrades discrimination, giving CRL + AC more room to improve.The state-only critic is already approximately 57% accurate on cube, so action-related gains are small in accuracy but important for policy selection.
Noise dependence: On puzzle, discrimination metrics remain nearly flat across noise because the state stays informative, yet action chunking still improves performance.This separates the performance benefit from noise-driven changes in state informativeness.
Network depth: CRL + AC matches or exceeds CRL at every tested network depth on Ant Hardest Maze and Humanoid.The comparison sweeps depths of 2, 4, 8, and 16 layers with H=3.
Network depth: A 2-layer CRL + AC network is comparable to a 16-layer CRL network on Ant Hardest Maze, indicating a large compute saving.The 16-layer model is 8× deeper than the 2-layer model.

E.4 REPLANNING INTERVAL ABLATION

The replanning ablation contrasts executing whole action chunks with replanning after every action. Replanning has little aggregate effect offline but substantially harms online performance, leaving reactive execution with the improved chunked critic unresolved.

Offline setting: Replanning after every action has no significant effect on aggregated offline success rate, although effects vary across environments.Offline comparisons use Hexec=1 versus executing the whole chunk with Hexec=H.
Online setting: Online replanning after every action substantially degrades performance, reducing it below standard CRL.The online ablation shortens replanning only during evaluation; experience collection retains Hexec=H.
Open questions: The offline–online contrast remains unexplained and may reflect differences in training protocols and evaluation environments.The authors leave a full explanation to future work.
Discount factor: The discount-factor ablation finds γ=0.95 better than γ=0.99 across all evaluated tasks for both CRL and CRL + AC.The paper therefore uses γ=0.95 as the default.
Other methods: Preliminary results show action chunking raises success rate by ∼30% for TD-JEPA and ∼50% for Forward-Backward.The source does not isolate whether these gains arise from representational effects or multi-step returns.

F.1 OGBENCH BEST FQL

Across offline OGBench environments, CRL with action chunking consistently outperforms standard CRL, including under replanning. The evaluation selects the best FQL regularization strength per environment and method.

OGBench results: CRL + AC consistently outperforms CRL across 21 offline OGBench environments, with an average improvement of 32%.Learning curves cover H∈{1,3,5} with 95% bootstrapped confidence intervals.
Replanning: CRL + AC with replanning outperforms CRL across many OGBench environments, with an average improvement of ∼30%.The replanning interval is Hexec=1.
FQL selection: The main experiments select the best-performing FQL α∈{1,3,10} separately for each environment and method.Additional curves report all evaluated α values.
Online comparison: Across 11 online JaxGCRL environments, CRL + AC with H=3 achieves approximately 90% average improvement over standard CRL on success rate.The environments span locomotion and navigation tasks.
Online comparison: For time at goal, CRL + AC with H=3 achieves +111.2% improvement over standard CRL on average across the online environments.Time at goal rewards task success and behavioral stability.
Replanning: In online tasks, single-step replanning underperforms open-loop chunk execution, unlike the offline setting.The online comparison uses Hexec=1 versus open-loop execution and is reported for success rate and time at goal.