PaperlayerView exact revision on arXiv

Source-linked AI summary

Behavior Generation with Latent Actions

Seungjae Lee, Yibin Wang, Haritheja Etukuru, H. Jin Kim, Nur Muhammad Mahi Shafiullah, Lerrel Pinto

arXiv:2403.03181v2cs.LGcs.AIcs.RO

TL;DR

Behavior generation must model multimodal continuous actions whose sequential errors can compound, while k-means tokenization struggles with high-dimensional and temporally extended actions. VQ-BeT uses transformer modeling with hierarchical residual vector quantization for action tokens, achieving strong results across behavior-generation settings and faster diffusion-model inference. The paper reports improvements in unconditional and conditional benchmarks, real-world long-horizon manipulation, autonomous driving, and inference speed.

  • Problem

    Behavior generation requires modeling multimodal continuous actions with temporal dependencies, but k-means tokenization does not scale well to high-dimensional or temporally extended actions.

  • Method

    VQ-BeT combines transformer-based long-horizon modeling with residual vector quantization to learn discrete latent representations of actions or action chunks.

  • Results

    VQ-BeT achieves SOTA unconditional performance in 5/7 environments, conditional performance in 6/7, and improves real-world long-horizon manipulation by 73%.

  • Takeaways & Limitations

    VQ-BeT provides a versatile behavior-generation model spanning conditional and unconditional tasks, simulated manipulation, autonomous driving, and real-world robotics.

  • Takeaways & Limitations

    The method is motivated by limitations in BeT and C-BeT tokenizers, including sensitivity to k-means cluster count and difficulty capturing high-dimensional correlated action distributions.

Abstract

from arXiv · show

Generative modeling of complex behaviors from labeled datasets has been a longstanding problem in decision making. Unlike language or image generation, decision making requires modeling actions - continuous-valued vectors that are multimodal in their distribution, potentially drawn from uncurated sources, where generation errors can compound in sequential prediction. A recent class of models called Behavior Transformers (BeT) addresses this by discretizing actions using k-means clustering to capture different modes. However, k-means struggles to scale for high-dimensional action spaces or long sequences, and lacks gradient information, and thus BeT suffers in modeling long-range actions. In this work, we present Vector-Quantized Behavior Transformer (VQ-BeT), a versatile model for behavior generation that handles multimodal action prediction, conditional generation, and partial observations. VQ-BeT augments BeT by tokenizing continuous actions with a hierarchical vector quantization module. Across seven environments including simulated manipulation, autonomous driving, and robotics, VQ-BeT improves on state-of-the-art models such as BeT and Diffusion Policies. Importantly, we demonstrate VQ-BeT's improved ability to capture behavior modes while accelerating inference speed 5x over Diffusion Policies. Videos and code can be found https://sjlee.cc/vq-bet

1. Introduction

VQ-BeT targets multimodal, temporally correlated behavior generation by combining transformer sequence modeling with vector-quantized action representations. Across simulated, driving, and robotic settings, it reports strong performance and faster inference than relevant baselines.

  • Motivation: Behavior generation is difficult because actions are multimodal, temporally correlated, noisy, and physically grounded, so small errors can cascade into failures.These challenges distinguish action generation from generation in language, images, audio, and video.
  • Approach: VQ-BeT combines transformers for long-horizon modeling with vector quantization for expressive, lower-cost representations of continuous actions.The method is designed for conditional and unconditional generation from uncurated behavior datasets.
  • Results: 5/7 environments: VQ-BeT achieves SOTA unconditional behavior generation, outperforming BC, BeT, and diffusion policies.Entropy metrics and qualitative visualizations indicate better capture of multiple behavior modes.
  • Results: 6/7 environments: VQ-BeT achieves SOTA conditional behavior generation, improving upon GCBC, C-BeT, and BESO by specifying goals as input.The comparison covers goal-conditional behavior generation across seven environments.
  • Efficiency: 5× speedup in simulation and 25× on real-world robots: VQ-BeT is faster than diffusion-based multi-pass models.Its single-pass design generates an action or action chunk in one forward pass.
  • Results: 73% improvement on long-horizon tasks: VQ-BeT scales to real-world robotic manipulation including pick-and-place and door closing.The paper also reports direct applicability to autonomous-driving benchmarks such as nuScenes.

2. Background and Preliminaries

Behavior cloning must represent multimodal actions, temporal context, partial observability, and optional goals. BeT uses discrete bins plus continuous offsets, while VQ-BeT replaces k-means-style tokenization with learned vector quantization that scales to complex action spaces.

  • Behavior cloning: Behavior cloning learns a mapping from observations to continuous actions, but multimodality motivates alternative losses and distributions conditioned on history or goals.The modeled distributions may use observation histories and explicit or implicit goal variables.
  • Behavior Transformers: BeT and C-BeT use GPT-like transformers to model temporal correlations and multiple behavior modes through discrete components with continuous offsets.Their tokenization lets each learned mode remain continuous after selecting a categorical component.
  • Behavior Transformers: k-means tokenization is sensitive to the choice of k and may merge modes, split modes, or miss structure in high-dimensional correlated action chunks.These limitations motivate a learned discrete latent embedding space for actions and action chunks.
  • Vector Quantization: Vector quantization maps a continuous input x to the nearest embedding vector zq in a finite codebook.The discretization bottleneck replaces continuous values with learned codebook vectors.
  • Residual Vector Quantization: Residual VQ uses multiple quantization stages whose codebook vectors are summed, successively quantizing residuals to represent embeddings.This multi-stage construction replaces each vanilla VQ-VAE embedding with a sum drawn from finite codebooks.

q. The residual, x −z1

The input is iteratively quantized through successive layers that pass updated residuals forward. The final quantized vector sums vectors from finite codebooks.

  • Successive quantizing layers update and pass the residual x − Pp_q to the next layer.
  • The quantization process uses Nq − 1 quantizing layers before producing the final quantized input vector.
  • The final quantized input vector is constructed as a sum of vectors from finite codebooks.

3. Vector-Quantized Behavior Transformers

VQ-BeT learns discrete latent representations for continuous action sequences with Residual VQ, then predicts hierarchical action codes using a transformer. The model supports both unconditional and goal-conditional behavior generation.

  • Architecture: VQ-BeT uses a two-stage pipeline: action discretization with Residual VQ followed by transformer-based VQ-BeT learning.Residual VQ encodes action or action chunks into discrete latent codes, which the transformer predicts from observation sequences.
  • Action discretization: Residual VQ replaces k-means binning by learning a discrete latent embedding space for actions or action chunks.The learned representation is intended to scale to complex action spaces while preserving continuous-action fidelity through decoding.
  • Action discretization: Hierarchical residual quantization represents each action latent as a sum of vectors from multiple codebooks, with primary codes capturing coarse structure and secondary codes refining actions.The paper reports using two residual layers in its experiments.
  • Code prediction: The transformer predicts categorical codes from observation histories, while an offset head adjusts discretized action centers to maintain fidelity.The predicted residual embeddings are summed and passed through a decoder to obtain the quantized behavior.
  • Task formulation: VQ-BeT supports unconditional policies conditioned on observation histories and conditional policies additionally conditioned on current and future goal observations.The same architecture is described as applicable to both conditional and unconditional generation, with an optional goal input.

4. Experiments

The experiments examine VQ-BeT in conditional and unconditional behavior generation, focusing on task performance and diversity. They specifically test policy quality and the model’s ability to capture multimodal behavior.

  • Evaluation goals: Experiments evaluate VQ-BeT on both conditional and unconditional behavior-generation tasks across different environments and datasets.The study asks how well VQ-BeT policies perform in each setting.
  • Evaluation goals: Task quality is measured by how well generated behavior achieves an objective or goal, while diversity is measured by entropy over accomplished subtasks or goals.These metrics are used to assess both performance and behavioral variation.
  • Evaluation goals: The experiments directly assess whether VQ-BeT captures the multimodality present in the behavior dataset.This is stated as a separate experimental question alongside environment-level policy performance.

3. Does VQ-BeT scale beyond simulated tasks?

VQ-BeT extends beyond simulated behavior generation to autonomous driving and real-world robotic manipulation while retaining competitive performance. Its single-pass generation also reduces inference cost relative to diffusion-based multi-pass models.

  • Behavior diversity: VQ-BeT captures both behavior modes in the PushT dataset and executes both possible door-closing orders in the “Close Fridge and Toaster” task.These rollouts provide qualitative evidence of multimodal behavior rather than collapse to a single mode.
  • Autonomous driving: VQ-BeT matches or exceeds task-specific state-of-the-art methods on the nuScenes autonomous-driving benchmark.It achieves the lowest average L2 trajectory-following error and collision probability that is better or on-par with older self-driving methods, despite not being designed specifically for driving.
  • Design choices: Residual vector quantization is important for expressivity, while autoregressive VQ-code prediction negatively affects performance in the kitchen environment.Removing residual VQ or equally weighting the two VQ layers causes a significant performance drop.
  • Real-world robotics: VQ-BeT outperforms all baselines on real-world two-phase manipulation tasks by a relative margin of 73%.On single-phase tasks, both VQ-BeT and DiffusionPolicy-T achieve at least 90% success, while the longer-horizon comparison favors VQ-BeT.
  • Real-world robotics: VQ-BeT achieves at least three times DiffusionPolicy’s success rate at the end of four long-horizon real-world tasks.The experiments involve sequences of three or more subtasks across communal-kitchen and conference-room scenes with varied objects.

5. Related Works

Related work develops behavior generation through generative modeling, action reparameterization, and transformer-based modeling of multimodal or partially observed behaviors. VQ-BeT follows this trajectory by learning discrete latent action representations rather than relying on fixed action-space assumptions.

  • Deep generative models for modeling behavior: Generative behavior modeling spans inverse reinforcement learning, multimodal action-distribution fitting, and generalized learning from large behavior datasets.The cited literature applies generative-modeling tools to learn diverse behaviors from human-collected and multimodal datasets.
  • Action reparametrization: Action reparameterization replaces or restructures continuous actions, but independent, causal, handcrafted, and temporally abstract representations each impose scaling or modeling constraints.The literature includes bang-bang discretization, dimension-wise discretization, learned reparameterization, options, and handcrafted primitives.
  • Action reparametrization: Learned action reparameterization avoids committing to independent or causally dependent action dimensions when no prior action-space structure is available.This motivates learned representations such as those used by VQ-BeT.

6. Conclusion

VQ-BeT learns behavior from open-ended, multimodal data by tokenizing actions with a residual VQ-VAE and predicting action tokens with a transformer. The paper demonstrates performance across manipulation, locomotion, and self-driving tasks, while identifying larger-scale behavior datasets and cross-embodiment transfer as future directions.

  • VQ-BeT tokenizes continuous actions using a residual VQ-VAE and predicts the resulting action tokens with a transformer.
  • The model performs well across manipulation, locomotion, and self-driving tasks.
  • Scaling VQ-BeT to behavior datasets with substantially more data, environments, and behavior modes remains an open direction.
  • A shared action latent space could support policy translation between different robots or between humans and robots.
  • The authors identify faster real-world reinforcement learning using a learned discrete action space as future work.

A. Experimental and Dataset

The experiments cover simulated manipulation, locomotion, self-driving, and real-world robotic environments with varied action spaces, goals, observations, and multimodal demonstrations. These settings test both conditional and unconditional behavior generation across diverse task structures.

  • The evaluation spans Franka Kitchen, PushT, BlockPush, UR3 BlockPush, Multimodal Ant, nuScenes self-driving, and a real-world kitchen environment.
  • Simulated environments: Franka Kitchen uses a 7-dimensional action space, 566 human demonstrations, seven possible tasks, and state- or image-based observations.
  • Simulated environments: PushT evaluates two-dimensional end-effector velocity control for moving a T-shaped block, including image-based and goal-conditioned variants.
  • Simulated environments: BlockPush and UR3 BlockPush test multimodal block-moving behavior with alternative target assignments or block orders.
  • Simulated environments: Multimodal Ant requires reaching corner goals in different orders, with conditional and unconditional evaluations using two or four goals.
  • Beyond manipulation: The nuScenes setup uses 684 demonstration scenes and predicts the next six driving timesteps from object-centric data with rollout-time direction goals.
  • Real-world environments: The real-world kitchen environment randomizes gripper and can positions, uses 45 demonstrations per task, and includes multimodal two-task sequences.

B. Additional Results

Additional evaluations compare VQ-BeT with related baselines across conditional and unconditional tasks, multimodal trajectory generation, action-code structure, and real-robot control. The figures and tables emphasize broad benchmark performance, smooth multimodal trajectories, and inference-time comparisons.

  • Self-driving evaluation: Table 10 evaluates nuScenes trajectory planning with lower-is-better metrics and notes that Agent-Driver uses additional road-lane and shoulder information.
  • Multimodal behavior: VQ-BeT generates smooth trajectories for both T-block-pushing modes with short single-step inference time.
  • Real-world evaluation: Table 11 compares diffusion policy under closed-loop and receding-horizon control in real-world robot experiments.
  • Action representations: Primary action codes cluster coarse ranges, while all 144 combinations of full codes provide finer decoded action structure.
  • Simulation benchmarks: VQ-BeT achieves the best performance in most conditional simulation environments and matches the best baseline on BlockPush.
  • Simulation benchmarks: VQ-BeT achieves the best performance in most unconditional simulation environments and matches the best baseline on BlockPush and Image Kitchen.

B.1. VQ-BeT with larger Residual VQ Codebook

The paper tests VQ-BeT with substantially larger residual VQ codebooks and optional masking of code combinations absent from the dataset. Performance is generally robust to codebook expansion, with environment-specific trade-offs and no consistent masking advantage.

  • Increasing code combinations by ×10 to ×250 has little impact on performance in most environments.
  • Performance slightly increases with larger codebooks for conditional Ant Multimodal and unconditional PushT.
  • Larger codebooks create performance–entropy trade-offs in unconditional Ant Multimodal and Kitchen.
  • Kitchen Conditional is the only environment where performance decreases with the extended codebook size.
  • Masking absent code combinations has no consistent benefit across environments.
  • VQ-BeT remains robust when the codebook is large enough to capture the dataset’s major modes.

C.2. VQ-BeT for Driving Dataset

For the NuScenes driving dataset, VQ-BeT is adapted to variable numbers of detected objects by using specialized tokens for mission, ego state, trajectory history, and objects.

  • NuScenes scenes can contain different numbers of detected objects, creating variable observation dimensionality for VQ-BeT.The model input types are modified to accommodate this change.
  • Input tokens: Mission tokens encode whether the agent should go forward, turn left, or turn right.
  • Input tokens: Ego-state tokens contain velocity, angular velocity, acceleration, heading speed, and steering angle.
  • Input tokens: Trajectory history tokens contain the ego's trajectories and velocities from the last 2 seconds.
  • Input tokens: Object tokens encode each object's current position, predicted future position, perception and prediction outputs, and one of 15 object classes.
Loading 2403.03181v2…