Efficient Planning in a Compact Latent Action Space

TL;DR

The paper tackles the challenge of planning in high-dimensional continuous action spaces by learning a compact, discrete latent action space with a state-conditioned VQ-VAE. TAP uses a Transformer-based prior to sample latent codes that reconstruct long-horizon trajectories, enabling efficient planning via beam search and returning high-quality plans in offline RL settings. Key findings show TAP achieves strong performance across locomotion, Adroit, and AntMaze benchmarks, with significantly lower decision latency than baselines like Trajectory Transformer, especially as action dimensionality grows. The work demonstrates that planning in a learned latent space can deliver practical, real-time capable planning for high-dimensional control tasks, with robust ablations highlighting the importance of state conditioning, latent steps, and an OOD-aware objective.

Abstract

Planning-based reinforcement learning has shown strong performance in tasks in discrete and low-dimensional continuous action spaces. However, planning usually brings significant computational overhead for decision-making, and scaling such methods to high-dimensional action spaces remains challenging. To advance efficient planning for high-dimensional continuous control, we propose Trajectory Autoencoding Planner (TAP), which learns low-dimensional latent action codes with a state-conditional VQ-VAE. The decoder of the VQ-VAE thus serves as a novel dynamics model that takes latent actions and current state as input and reconstructs long-horizon trajectories. During inference time, given a starting state, TAP searches over discrete latent actions to find trajectories that have both high probability under the training distribution and high predicted cumulative reward. Empirical evaluation in the offline RL setting demonstrates low decision latency which is indifferent to the growing raw action dimensionality. For Adroit robotic hand manipulation tasks with high-dimensional continuous action space, TAP surpasses existing model-based methods by a large margin and also beats strong model-free actor-critic baselines.

Figures (8)

• Figure 1: (a) gives an overview of TAP modelling, where blocks represent the latent actions. (b) shows decision time growth with the dimensionality $D$. Tests are done on a single GPU. The number of planning steps for (b) is 15 and both models apply a beam search with a beam width of 64 and expansion factor of 4. (c) shows the relative performance between TAP and baselines when dealing with tasks with increasing raw action dimensionalities.
• Figure 2: Illustration of the training and test time inference process of TAP. The left-hand side shows the training process, highlighting the design of the bottleneck. The right-hand side figure shows how we generate plans during the test time, with vanilla sampling.
• Figure 3: Results of ablation studies, where the height of the bar is the mean normalised scores on gym locomotion control tasks.
• Figure 4: (a) Illustration of dimension-wise autoregressive modelling used by Trajectory Transformer. Blocks represent discretised state/action dimensions or discrete latent variables. (b) shows TAP style modelling. (c) Horizon Scalability, tested for a low-dimensional task, hopper.
• Figure 5: Visualization of different antmaze environments.