[Paper Review] Action-Quantized Offline Reinforcement Learning for Robotic Skill Learning

Jianlan Luo, Perry Dong+4 more|arXiv (Cornell University)|Oct 18, 2023

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes State-Conditioned Action Quantization (SAQ), a VQ-VAE-based method that discretizes continuous action spaces in offline reinforcement learning to enable exact computation of policy constraints and conservatism regularizers. SAQ improves performance by 2–3× on long-horizon robotic manipulation tasks in Robomimic, outperforming continuous-action counterparts and prior offline RL methods on narrow expert demonstration datasets.

ABSTRACT

The offline reinforcement learning (RL) paradigm provides a general recipe to convert static behavior datasets into policies that can perform better than the policy that collected the data. While policy constraints, conservatism, and other methods for mitigating distributional shifts have made offline reinforcement learning more effective, the continuous action setting often necessitates various approximations for applying these techniques. Many of these challenges are greatly alleviated in discrete action settings, where offline RL constraints and regularizers can often be computed more precisely or even exactly. In this paper, we propose an adaptive scheme for action quantization. We use a VQ-VAE to learn state-conditioned action quantization, avoiding the exponential blowup that comes with naïve discretization of the action space. We show that several state-of-the-art offline RL methods such as IQL, CQL, and BRAC improve in performance on benchmarks when combined with our proposed discretization scheme. We further validate our approach on a set of challenging long-horizon complex robotic manipulation tasks in the Robomimic environment, where our discretized offline RL algorithms are able to improve upon their continuous counterparts by 2-3x. Our project page is at https://saqrl.github.io/

Motivation & Objective

• To address distributional shift and overestimation in continuous-action offline RL by enabling exact implementation of policy constraints and conservatism regularizers.
• To overcome the exponential complexity of naive action discretization while preserving precision through adaptive, state-dependent quantization.
• To improve sample efficiency and performance on narrow, expert demonstration datasets common in robotic learning.
• To simplify and stabilize offline RL training by replacing continuous-action approximations with exact discrete-action computations.
• To validate the method on challenging long-horizon robotic manipulation tasks where continuous offline RL fails.

Proposed method

• Uses a VQ-VAE to learn a state-conditioned codebook that maps states to discrete action embeddings, enabling adaptive action quantization.
• Performs action discretization in a state-dependent manner, learning only the most in-distribution action primitives relevant to each state.
• Applies the discretized action space to three state-of-the-art offline RL algorithms: IQL, CQL, and BRAC, enabling exact computation of policy constraints and conservative value functions.
• Trains the VQ-VAE on behavior demonstration data to learn a compact, task-relevant action codebook that avoids exponential blowup.
• Uses the learned codebook to replace continuous actions in the RL algorithm, allowing precise and stable training without approximation errors.
• Employs a two-stage training process: pre-train the VQ-VAE on demonstration data, then fine-tune the RL policy with the discrete action space.

Figure 1: Left : Training SAQ. Right : Discrete policy training/evaluation with SAQ. SAQ learns a scalar discrete representation of continuous actions using state-conditioned VQ-VAE. During policy training and evaluation time, we use the decoder of SAQ to reconstruct the continuous actions from the

Experimental results

Research questions

• RQ1 Can state-conditioned action quantization improve the performance of offline RL algorithms on narrow, expert demonstration datasets?
• RQ2 Does discrete action space enable exact computation of conservatism and policy constraints, reducing approximation errors in offline RL?
• RQ3 How does SAQ compare to continuous-action offline RL in long-horizon robotic manipulation tasks?
• RQ4 Can SAQ enable offline RL to surpass imitation learning baselines on challenging robotic environments like Robomimic?
• RQ5 What is the impact of adaptive, state-dependent discretization on sample efficiency and training stability?

Key findings

• SAQ improves performance by 2–3× on long-horizon robotic manipulation tasks in the Robomimic environment compared to continuous-action offline RL methods.
• On the Robomimic Proficient Human dataset, SAQ-IQL achieves 90.8% success rate on the lift task, outperforming the continuous IQL baseline of 58%.
• SAQ-CQL achieves 92.7% success on the lift task, significantly exceeding the continuous CQL result of 64.2%.
• On the transport task, SAQ-BC achieves 66.4% success, outperforming the continuous BC baseline of 31.73%.
• The average success rate across all Robomimic tasks improves from 41.71% with SAQ-BC to 64.12% with the original Robomimic BC baseline, indicating SAQ’s strong generalization.
• The method enables exact enforcement of conservatism and policy constraints, eliminating approximation errors common in continuous-action settings.

Figure 2: Left : visualization of the maze2d environment. Middle : the average return of SAQ-CQL and continuous CQL over the course of offline training. Right : estimated and exact CQL conservatism penalty values. The estimated conservative penalty for continuous-action CQL diverges significantly fr