PaperlayerView exact revision on arXiv

Source-linked AI summary

Behavior-Aligned Action Tokenization for Robot Policy Learning

Junbo Dong, Ze Chen, Zhendong Xie, Junjie Li, Lixin Xu, Xuemin Chi, Yiming Song, Zhaoyuan Ma

arXiv:2609.27513v1cs.RO

TL;DR

Different tasks can contain similar local motions with different timing, while existing tokenizers provide limited explicit behavioral correspondence supervision. BAAT uses Soft-DTW to select action-chunk pairs and aligns quantized coordinates jointly with reconstruction. It achieves the highest mean simulation success among compared methods, and a controlled LIBERO-All ablation reports higher policy success alongside lower replay success.

  • Problem

    Similar local motions recur across tasks with different timing, while reconstruction alone does not specify which action chunks across demonstrations should have nearby representations.

    Source: §I · p. 1

  • Method

    BAAT uses Soft-DTW to select corresponding action chunks and jointly aligns their quantized coordinates with reconstruction before training an autoregressive policy on the resulting targets.

    Source: §I · p. 1

  • Results

    BAAT achieves the highest mean simulation success among compared methods, and controlled LIBERO-All alignment raises policy success from 70.2% to 79.0% while replay success falls from 80.5% to 71.5%.

    Source: §I · pp. 1–2

  • Takeaways & Limitations

    The findings support behavioral correspondence as supervision for organizing shared motion structure in action tokenizers for multi-task policy learning.

    Source: §V · p. 8

Abstract

from arXiv · show

Autoregressive robot policies learn continuous control by predicting discrete action tokens from observations. Different tasks often share local motions, yet behavioral correspondence across demonstrations receives limited explicit supervision in existing tokenizers. Motions with different timing can therefore lack a shared representation despite following similar patterns. We propose Behavior-Aligned Action Tokenization (BAAT), which uses soft dynamic time warping (Soft-DTW) to select corresponding action chunks and aligns their quantized coordinates jointly with reconstruction. This objective encourages similar motions across tasks to occupy nearby quantized representations while retaining executable action detail. A history-conditioned diffusion decoder reconstructs continuous action chunks from these tokens, and a downstream autoregressive policy learns to predict them. We evaluate BAAT on selected tasks from three simulation benchmarks and two real robot tasks. BAAT achieves a mean simulation success rate of approximately 45.2%, exceeding OAT by approximately 7.2 percentage points. In the controlled LIBERO-All alignment ablation, policy success rises from 70.2% to 79.0% while trajectory replay success decreases. These results support behavioral correspondence as supervision for organizing shared motion structure in action tokenizers and improving downstream robot policy learning.

I. INTRODUCTION

Similar local motions recur across tasks with different timing, but reconstruction alone does not specify which action chunks should share nearby representations. BAAT uses Soft-DTW correspondence to align quantized action representations jointly with reconstruction, then evaluates the resulting tokenizer for policy learning.

I. INTRODUCTION

  • Chunk reconstruction alone does not specify which motions across demonstrations should occupy nearby representations, motivating behavioral correspondence across task boundaries.Similar local motions can recur with different timing, such as reaching, lifting, or lateral transfer.Source: §I · p. 1
  • The intended shared structure is a neighborhood of similar local motions, not shared task semantics or identical discrete codes.Alignment is applied after quantization while reconstruction retains executable action detail.Source: §I · p. 1
  • Soft-DTW selects corresponding action chunks, and behavioral alignment constrains their quantized coordinates jointly with action reconstruction.The tokenizer is then frozen to provide categorical targets for a downstream autoregressive policy.Source: §I · p. 1
  • In controlled LIBERO-All ablation, alignment raises policy success from 70.2% to 79.0% while replay success falls from 80.5% to 71.5%.Cross-task retrieval also shows corresponding local motions in BAAT’s quantized neighborhoods.Source: §I · p. 1
  • BAAT achieves the highest mean simulation success among compared methods across evaluations spanning three simulation benchmarks and two real robot tasks.The introduction also reports an approximately 8.1 percentage-point gain from joint LIBERO training.Source: §I · pp. 1–2

II. RELATED WORK

Prior action tokenizers organize policy targets through compression, token ordering, stability, or multimodal supervision. BAAT instead uses action-chunk motion correspondence and conditional denoising reconstruction to organize quantized action representations.

A. Action Tokenization for Robot Policies

  • FAST compresses temporally structured action targets with discrete cosine transform coefficients and byte-pair encoding, while OAT orders FSQ tokens for progressively longer-prefix reconstruction [4], [5].Direct coordinate binning makes sequence length grow with action dimension and prediction horizon [17].Source: A. Action Tokenization for Robot… · p. 2
  • ActionCodec studies token stability and uses adjacent chunks as temporal-contrastive positives [8], while X-Tokenizer adds hierarchical and foundation-model supervision [9].BAAT instead selects mini-batch action chunks by motion correspondence, without requiring temporal adjacency or task labels.Source: A. Action Tokenization for Robot… · p. 2

B. Behavior-Aware Action Representation

  • Act2Vec learns action embeddings from demonstration context, while Dynamics-Aware Embeddings jointly represent states and action sequences through forward prediction.Temporal correspondence methods include cycle-consistent video matching and action-centric recovery from video.Source: B. Behavior-Aware Action Representation · p. 2

C. Generative Action Modeling and Reconstruction

  • Diffusion Policy generates action trajectories through conditional denoising [24], and Octo integrates a diffusion action head into a generalist policy [25].BAAT uses conditional denoising to reconstruct continuous actions from behavior-aligned quantized representations and executed action history.Source: C. Generative Action Modeling and… · p. 2

III. METHOD

BAAT organizes similar local motions into shared quantized neighborhoods while preserving action detail for execution. Its training selects behaviorally corresponding chunks, aligns their quantized coordinates during reconstruction, and then trains a policy to predict the resulting categorical targets.

A. Method Overview

  • BAAT first selects action-chunk pairs with Soft-DTW and jointly trains an encoder and history-conditioned diffusion decoder using alignment and action reconstruction.It then freezes the tokenizer and decoder, trains an autoregressive policy on categorical labels, and decodes predicted labels with executed action history.Source: A. Method Overview · p. 2

1) Action representation:

  • Each action chunk is encoded into latent slots, projected to scalar coordinates, and quantized into categorical targets and corresponding quantized coordinates.FSQ quantizes each scalar independently; the policy predicts categorical targets, while alignment and decoding use the coordinates.Source: 1) Action representation · p. 2

2) Behavioral alignment:

  • Because similar motions can differ in speed or phase, Soft-DTW selects low-cost chunk pairs under temporal warping for behavioral alignment.The alignment loss acts on paired quantized representations, encouraging nearby coordinates without requiring identical codes; joint reconstruction penalizes loss of executable detail.Source: 2) Behavioral alignment · pp. 2–3, Figure 1 · p. 3

C. History-Conditioned Diffusion Reconstruction

A history-conditioned diffusion decoder reconstructs continuous action chunks from quantized representations while using preceding actions as execution context. The reconstruction objective masks invalid action elements, and the real-robot configuration adds a residual parameterization for arm targets.

C. History-Conditioned Diffusion Reconstruction

  • The decoder predicts a clean action chunk from its noisy version, diffusion step, quantized representation, and preceding-action history.The noisy chunk is formed using a noise schedule, and the decoder minimizes masked Smooth-L1 loss.Source: C. History-Conditioned Diffusion… · p. 3
  • Valid action elements are selected by a mask, and a straight-through gradient lets reconstruction and alignment losses update the encoder and projection through quantization.The joint tokenizer objective uses an alignment weight.Source: C. History-Conditioned Diffusion… · p. 3
  • For the evaluated 12D bimanual configuration, the immediately preceding action conditions the decoder through an MLP and FiLM modulation.The history action is normalized using action statistics.Source: C. History-Conditioned Diffusion… · p. 3
  • The real-robot residual predicts ten arm-joint offsets from a shared preceding-action anchor, while the two grippers are predicted as absolute targets.Adding the anchor recovers normalized absolute arm targets; this configuration also uses boundary and overlap losses and a longer prediction horizon.Source: C. History-Conditioned Diffusion… · p. 3

D. Factorized Autoregressive Policy Learning

The frozen tokenizer supplies discrete targets for an autoregressive policy conditioned on observations and an optional task input. At inference, predicted latent labels are converted to quantized coordinates that condition the decoder alongside executed action history.

D. Factorized Autoregressive Policy Learning

  • The policy is trained with teacher forcing to predict tokenizer-produced latent targets from observations and an optional task input.It autoregresses over latent slots and predicts each slot’s FSQ factors independently in parallel using a Km-class head per factor.Source: D. Factorized Autoregressive Policy… · pp. 3–4
  • At inference, previously generated latent slots condition the policy’s next predictions.The map r converts sampled FSQ factor labels to quantized coordinates for decoding.Source: D. Factorized Autoregressive Policy… · p. 4
  • The frozen diffusion decoder uses predicted quantized coordinates and executed action history to generate each next chunk by iterative denoising.The figure caption identifies example scenes from the LIBERO suites, RoboTwin 2.0 tasks, and RoboCasa tasks.Source: D. Factorized Autoregressive Policy… · p. 4, Figure 2 · p. 4

IV. EXPERIMENTS

The experiments assess cross-task behavioral neighborhoods, control and action fidelity across training regimes, and simulation ablations of alignment and reconstruction. They cover selected simulation tasks and two real-robot tasks.

IV. EXPERIMENTS

  • The experiments first inspect cross-task behavioral neighborhoods, then assess control and action fidelity across training regimes and simulation ablations.A separate offline decoder-history comparison diagnoses chunk continuity on recorded real-robot data.Source: §IV · p. 4

1) Tasks and training regimes:

  • The study evaluates 40 LIBERO tasks through per-suite training or joint LIBERO-All training across demonstrations from all four suites.It also evaluates four RoboCasa and four RoboTwin 2.0 tasks, each with joint training across selected tasks, plus toothbrush organization and object sorting on two bimanual systems.Source: 1) Tasks and training regimes · p. 4

2) Baselines and policy training:

  • Baselines include FAST, direct coordinate binning, OAT, and continuous-action Diffusion Policy.Discrete interfaces share the downstream autoregressive architecture and training procedure where applicable; BAAT uses 16 latent slots, four FSQ factors per slot, and a 20-step prediction horizon with the first 16 steps executed.Source: 2) Baselines and policy training · p. 4

3) Metrics and Evaluation Protocol:

The evaluation reports closed-loop policy success and action-fidelity measures, with simulation ablations varying alignment weight under otherwise identical settings. Cross-task retrieval identifies recurring local motion patterns, but does not isolate alignment’s role in forming them.

3) Metrics and Evaluation Protocol:

  • Policy success is measured by closed-loop task success rate, while simulation reconstruction fidelity is measured by MSE and replay success.Simulation comparisons and ablations average over three seeds and evaluate 50 rollouts per task at each run’s final checkpoint; real-robot results report success rate and successful trials out of 20 attempts per task and method.Source: 3) Metrics and Evaluation Protocol · p. 4
  • The alignment ablation varies only alignment weight in LIBERO-All, holding other training and evaluation settings identical.The separate offline diagnostic measures arm reconstruction MAE and boundary-jump quantiles on recorded real-robot data without robot execution.Source: 3) Metrics and Evaluation Protocol · p. 4

B. Cross-Task Behavioral Neighborhoods

  • Nearest-neighbor retrieval reveals lateral-transfer, forward-reach, and vertical-lift patterns across tasks despite differences in trajectories, objects, and scene layouts.The retrieval compares LIBERO-10 queries with neighbors from LIBERO-Goal, LIBERO-Object, and LIBERO-Spatial; the local correspondence does not require a shared complete manipulation goal.Source: B. Cross-Task Behavioral Neighborhoods · p. 4
  • The retrieval evidence is qualitative and neither isolates alignment’s role in forming neighborhoods nor establishes independently callable skills.Joint-training comparisons and alignment ablations examine complementary control benefits and fidelity trade-offs.Source: B. Cross-Task Behavioral Neighborhoods · p. 4

C. Downstream Policy Performance

BAAT has the highest mean policy success across the compared simulation benchmarks, with its clearest LIBERO advantage under joint training. Cross-task neighborhoods show similar local motions, while joint-training comparisons do not isolate alignment from changes in demonstration volume and diversity.

1) Simulation benchmarks:

  • Approximately 45.2% mean policy SR makes BAAT highest among compared methods on each simulation benchmark, exceeding OAT by approximately 7.2 percentage points.On LIBERO-All, BAAT reaches 79.0%, approximately 2.5 points above OAT.Source: 1) Simulation benchmarks · p. 4, Table I · p. 5
  • LIBERO-10 queries retrieve cross-task chunks with similar lateral-transfer, forward-reach, and vertical-lift patterns, despite trajectory differences.These retrieved trajectories are shown without decoder generation or selection from trajectory plots.Source: Figure 3 · p. 5, 2) Per-suite versus joint training · p. 5

2) Per-suite versus joint training:

  • BAAT gains most from joint LIBERO training: its mean policy SR rises from 70.9% per-suite to 79.0% jointly, an approximately 8.1-point gain.Its per-suite average is below OAT and DP, and it trails FAST, OAT, and DP on independently trained LIBERO-10; FAST, OAT, and Bin gain 1.6, 1.3, and 2.5 points, respectively.Source: 2) Per-suite versus joint training · p. 5, Table II · p. 5
  • Joint training changes demonstration volume and diversity, so its larger gain does not isolate alignment’s contribution.The comparisons support the design motivation of organizing shared local motions, rather than establishing that alignment alone caused the gain.Source: 2) Per-suite versus joint training · p. 5
  • Reconstruction MSE and replay SR evaluate only part of the deployed interface, because closed-loop control also requires predicting tokens from observations and acting on resulting states.These metrics do not directly measure whether BAAT’s tokens are easier to predict.Source: Table III · p. 5, 3) Action Reconstruction · p. 6

4) Real robot manipulation:

On real robots, BAAT records the highest observed success on toothbrush organization and ties DP on object sorting. A separate offline reconstruction comparison finds improved arm fidelity and chunk-boundary continuity with history, but its combined settings do not isolate history’s effect.

4) Real robot manipulation:

  • Across 20 trials per method and task, BAAT reaches 14/20 successes on toothbrush organization and ties DP at 9/20 on object sorting.All methods have lower success on the 56-DoF sorting task, whose hardware and manipulation demands also differ, so the difference cannot be attributed to action dimension alone.Source: 4) Real robot manipulation · p. 6, Table IV · p. 6
  • Real-robot deployment showed abrupt chunk transitions despite smooth within-chunk motion, motivating the offline continuity comparison.Pronounced transition jitter was absent in the simulation evaluations.Source: 4) Real robot manipulation · p. 6

5) Policy optimization:

  • In individual runs, BAAT exceeds FAST and DP at early checkpoints, while OAT remains stronger at several early and intermediate checkpoints; the curves do not establish generally faster convergence.The three-seed final-checkpoint comparison supports BAAT’s benefit in attained policy performance.Source: 5) Policy optimization · p. 6, Figure 5 · p. 6

D. Reconstruction With and Without Decoder History

  • The offline comparison uses 100 recorded bimanual episodes at 30 Hz, 12D absolute joint targets, and 20,000 training steps for each configuration.It uses 16 FSQ latent slots and ten training trajectories; history is the ground-truth action immediately before each window, without feedback from prior predictions or robot execution.Source: D. Reconstruction With and Without… · p. 6
  • The reconstruction comparison samples complete windows every 16 steps from frame 16 and decodes them with 27 deterministic diffusion steps.Quantiles pool valid consecutive-chunk boundaries; arm MAE averages elementwise absolute reconstruction errors over prediction windows.Source: D. Reconstruction With and Without… · pp. 6–7
  • With decoder history, boundary-jump p95 falls approximately 49.6% and arm MAE falls 43.9%, while gripper MAE does not improve.Boundary p95 moves closer to the ground-truth reference (2.926/2.940); within-chunk one-step jump p95 changes less, from 2.760 to 2.679, indicating the continuity improvement is concentrated at chunk joins.Source: D. Reconstruction With and Without… · p. 7, Table V · p. 7
  • The offline gains come from training trajectories with ground-truth history, and auxiliary losses, prediction horizons, and window coverage also differ between configurations.The results therefore characterize the combined configuration, not continuity under predicted history or closed-loop performance.Source: D. Reconstruction With and Without… · p. 7

E. Ablation Studies

The ablations test alignment strength, temporal matching, and decoder choice, revealing trade-offs between policy success, replay success, and reconstruction. Decoder comparisons cover jointly trained tokenizer configurations and do not isolate diffusion as the source of policy gains.

Alignment strength

  • At alignment weight 0.1, policy SR rises 8.8 percentage points to 79.0% while replay SR falls 9.0 points; both metrics fall to zero at 0.5.Among tested weights, 0.1 yields the highest policy SR; the deterioration at stronger weights is consistent with conflict between alignment and execution detail, but does not diagnose representation collapse.Source: Table VI · p. 7, 1) How does alignment strength affect… · p. 7

Temporal matching

  • Soft-DTW improves policy and replay SR over pointwise matching by 6.0 and 6.9 percentage points, respectively, though replay SR remains below the no-alignment reference.The gains are consistent with temporal matching accommodating timing offsets, but the comparison does not directly measure pair quality.Source: Table VII · p. 7, 2) Does temporal matching improve… · p. 7

Decoder choice

  • Diffusion achieves the highest policy SR at 79.0%, leading Transformer and flow matching by 9.6 and 4.3 percentage points, respectively.Transformer minimizes reconstruction MSE and flow matching maximizes replay SR; because each configuration is jointly trained, this comparison does not isolate a diffusion property as the source of policy gain.Source: Table VIII · p. 7, 3) How does decoder choice affect… · pp. 7–8

V. CONCLUSION

BAAT uses behavioral correspondence to organize shared motion structure for multi-task policy learning, and its evaluations show strong results across simulation and real robot tasks. The authors identify VLM-backbone pretraining with behavior-aligned tokens as a future direction.

V. CONCLUSION

  • BAAT uses behavioral correspondence to organize shared motion structure for multi-task policy learning, aligning quantized actions with Soft-DTW and reconstruction through a history-conditioned diffusion decoder.Cross-task retrieval reveals neighborhoods of similar local motions; moderate alignment improves policy success despite reduced replay success, while excessive alignment degrades both.Source: §V · p. 8
  • Across three simulation benchmarks and two real robot tasks, BAAT achieves the highest mean simulation success and largest joint LIBERO training gain among compared methods.The paper concludes that behavioral correspondence is a useful source of supervision for action tokenization.Source: §V · p. 8
  • Future work will test behavior-aligned tokens as VLM-backbone pretraining targets while blocking gradients from a continuous action expert.The proposed test asks whether shared motion structure in token targets improves representation learning and generalization in VLA models.Source: §V · p. 8
Loading 2609.27513v1…