Quantization-Free Autoregressive Action Transformer
Authors: Ziyad Sheebaelhamd
Michael Tschannen
Michael Muehlebach
Claire Vernade
Venue: NeurIPS 2025 spotlight ()
Keywords: Imitation Learning
Reinforcement Learning
Transformers
Primary area: reinforcement_learning
TL;DR: We introduce a quantization-free way to train autoregressive transformers for continuous action decision making, improving on discretized action methods.

Abstract: Current transformer-based imitation learning approaches introduce discrete action representations and train an autoregressive transformer decoder on the resulting latent code. However, the initial quantization breaks the continuous structure of the action space thereby limiting the capabilities of the generative model. We propose a quantization-free method instead that leverages Generative Infinite-Vocabulary Transformers (GIVT) as a direct, continuous policy parametrization for autoregressive transformers. This simplifies the imitation learning pipeline while achieving state-of-the-art performance on a variety of popular simulated robotics tasks. We enhance our policy roll-outs by carefully studying sampling algorithms, further improving the results.

Quantization-Free e Action Transformer
Ziyad Sheebaelhamd University of Tübingen
Michael Tschannen Google DeepMind
Michael Muehlebach
Max Planck Institute for Intelligent Systems
Claire Vermade
University of Tübingen
Abstract
Current transformer-based imitation learning approaches introduce discrete action representations and train an autoregressive transformer decoder on the resulting latent code. However, the initial quantization breaks the continuous structure of the action space thereby limiting the capabilities the generative model. We propose quantization-fiee method instead that leverages Generative Infinite- Vocabulary Transformers (GIVT) as a direct, continuous policy parametrization for autoregressive transformers. This simplifies the imitation learning pipeline while achieving state-of-the-ar art performance on variety of popular simulated robotics tasks. We enhance our policy roll-outs by carefully studying sampling algorithms, further improving the results.
1 Introduction
Generative modeling lies at the heart of modern machine learning, enabling the production of outputs that adhere user-specified objectives and constraints by navigating high-dimensional spaces of potential outcomes. Fundamentally, this process can be understood as form of search where naive or exhaustive strategies are computationally intractable due to the size of the space. We consider imitation learning for robotic systems as an extreme example, where the search space has uncountable cardinality and complex structure arising from nonlinear dynamics, interactions with the environment and the fundamental multimodality in human or robot behavior. In many real-world robotics applications, the control signals involved are continuous in nature [33
59 28 Current
state-of-fe of-the-art methods for modeling policies in continuous control tasks broadly fall into two which will
categories: autoregressive transformer models 26 45 and diffusion models
be discussed in the next two paragraphs.
Existing autoregressive policies, the hand, ignore the challenge of learning in continuous domain by discretizing actions 266 45). This discretization can introduce several drawbacks:
discards the inherent structure of the continuous space, increases complexity by adding a separate quantization step, and may limit expressiveness or accuracy when fine-grained control is required. Furthermore, the popular approach of learning VQ-VAE (52) of the action space involves nondiffer- entiable operations and therefore requires advanced tricks during representation learning to sidestep optimization difficulties and instabilities (24. 20 [8] 31]. Hence, our goal is to design autoregressive transformers that preserve the continuous nature of the action space, leveraging the smooth and potentially multimodal behavior patterns that experts exhibit.
Diffusion policies, on the other hand, are adept at capturing continuous action distributions, but face two critical challenges in control tasks. First, real-time applications like robotics demand fast, efficient inference, yet diffusion models rely on iterative sampling, requiring dozens of forward passes
Correspondence: Ziyad Sheebaelhamd <ziyadd.shebbeelhnndd uni www.andddidduuuuunnene.ee.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

M-Mixture Parameters
Action Distribttion
Mistar PrttNiiiis
Mone
R.
(Cipomma))
Sequease
Decoder-Only former
Linear
Aud
Supartce
(a) An overview of the Q-FAT architecture.
FT
LTI
(b) Sample trajectories generated by a Q-FAT policy.
Figure I: Figure (a) shows the overview of QFAT. A sequence ht previous states and h goal states are projected into the Transformer's embedding dimension using linear layer. The Transformer then predicts the action distribution by predicting the GMM means, variances and mixture probabilities. Figure (b) shows sample of generated trajectories from 3 representative environments, demonstrating the captured multi multi- modality in the sequence of solving tasks (Kitchen and UR3 Block Push) or the direction from which an object is approached (PushT).
to generate actions. This makes them orders-of- -foorstuu- slower than autoregressive or feedforward policies 26] Second, the downstream task of imitation learning often reinforcement learning 43
39
58 17 which typically requires exact action likelihoods for tasks like exploration and
off off-policy corrections
49
27 56]. Diffusion models, however, generate samples by integrating
(stochastic) differential equations 48 47
making likelihood estimation computationally expensive.
To address these limitations, we propose the use of Generative Infinite Vocabulary Transformer (GIVT) (50) that models the policy as a Gaussian Mixture Model (GMM) on top of a decoder-only transformer architecture [53] 41). This approach enables modeling continuous actions "tokens" by directly paramcterizing the components the GMM in each time step, avoiding the need for discretization and thus preserving the inherent geometry of the action space. This choice allows us to exploit the benefits of large-scale sequence modeling, such as autoregressive policy generation and access to explicit likelihood estimates, while avoiding the pitfalls of artificially segmenting continuous control signals into discrete bins. In addition, we discuss two different sampling strategies that reduce variability in the generated trajectories.
In summary, our work:
Eliminates the need for action quantization in autoregressive policy parameterizatton, sim- plifying behavioral cloning pipelines while retaining the ability to directly evaluate action likelihoods.
Proposes two sampling algorithms to reduce the variance of the generated trajectories, including global mode-finding algorithm.
Demonstrates state-of-the-art performance standard benchmarks, both for conditional and unconditional policy generation using proprioceptive and image-based states, confirming the effectiveness of our method for multimodal policy modeling.
The implementation is available at
2 Sequence to Action Prediction
We consider the following standard setting for imitation learning of simulated robotics tasks. agent equipped with continuous controls A evolves in continuous state space S in discrete time. In

---

behavioral cloning, we assume that so-called behavior policy &a collected a dataset D = { { N-11 of N trajectories in the environment, where each trajectory т = {so. ab...., -1 s of states si S and actions of A, with n' denoting the length of trajectory i. Note that do not impose a Markovian assumption on the behavior policy. In general aj might depend on a history of states {s} - where h is a fixed window.
is sequence
Our objective is to learn sequence-to-action predictor, or imitating policy we. We treat the objective supervised learning problem, where the parameter of our policy optimized minimize the negative logllikelihood with respect to To. E, 1- log( To do so, we minimize:
()) = N
Σ 1=0
Σ
log To (a; s(
Note that a standard result in variational inference [5] shows that this objective is equivalent to minimizing the Kullback- Leibler (KL) divergence between the behavioral policy πτ and the learned policy no. implying that the learned policy Ty should match, or clone, the behavior of тp. However, the choice of parametrization is erucial to guarantee good approximation of the behavior policy.
3 The Q-FAT Algorithm
Current transformer-based policy parametrization approaches [26 45 32 consist of first discretizing the action space, c.g. using a VQ- VAE [52 57]. to comply with the standard transformer architecture requiring discrete vocabulary. Inspired by the success of GIVT [50] in the image generation domain, our model Quantization-FFie Autoregressive Action Transformer (Q-FAT) adapts GIVT to the continuous action behavioral cloning setting, modeling the action distribution at each step by predicting the parameters of a multivariate Gaussian mixture model (GMM). Namely, Q-FAT models the probability of an action 0₁ E R at step as
(at 8t-1, 8t k))
nt M(r n n!) Σ
where TH. ui; and of are produced by the transformer decoder [53 at step t, conditioned on a history of previous states {s:)i=/_ with 8; E Rd. Using a state history of length h² instead of only using the most recent state accounts for the potential non-Markovian structure in the demonstration dataset. Rather than predicting the full covariance Σ per mixture component, we assume a diagonal Gaussian distribution and only predict the diagonal entries of the covariance matrix af While this simplifying assumption enables efficient likelihood computations, it does not impose independence across action dimensions in general. Q-FAT trained using standard teacher forcing with a causal attention mask. We highlight that we only use state feedback and do not condition on the predicted actions. We discuss the important architecture decisions Appendix D
3.1 Goal Conditioning
For some tasks such as tracking, a future goal can be appended to the input to condition the resulting policy. The conditioning formalized by defining set of goal states S S. We then learn goal-conditioned policy no (az 8t-1..... St-hr, gg) where g C G. In practice, we define as sequence of future states to reached, i.e., g = 1t where the goal states are typically generated from a high-level planner that plans instrumental goal states as done in Gupta et al. [1]. In our experiments, we simulate goal states from the demonstration dataset.
3.2 Action Sampling from a k-GMM
QFAT samples the next action from the conditional distribution, ag Te(-|s:-1,.... St-h+ (atk.))), g), which a GMM with components. In practice, the model can be trained to output a short sequence of actions that are then executed in open-loop fashion. This is done to ensure action consistency and avoid over-fitting to idle actions 77. For simplicity, we describe only the next-action generation setting.
3

---

One can directly sample from the GMM that parametrizes the policy. For the vanilla CMM sampling. the procedure is:
1. Sample a mixture index k from the categorical distribution defined by the probabilities [
2. Sample an action a from the corresponding Gaussian component N (N,,,
While this approach directly samples from the learned policy distribution, it can degrade performance in practice when the dataset is noisy. In that case, the GMM may overfit the data distribution and produce noisy actions, resulting in jitter or unstable trajectories. We discuss below how Q-FAT can be easily equipped with stabilizing sampling techniques that directly exploit the fitted GMM distribution.
3.3 Stabilizing Output Trajectories
When controlling real-world robotics systems, it is important to generate action sequences without jitter, high frequency inputs may damage actuators and excite undesirable structural modes (2]. Thanks to the richness of the GMM representation of the policy, we can directly design stabilization patches.
Down-scaling the Mixture Variances A common heuristic to mitigate noisy samples is to adjust the sampling temperature, which in GMM parametrization is equivalent down-scaling the variance each Gaussian component 29 16 50 51). Concretely, one replaces each variance of by of for small (0, sometimes by several orders of magnitude (a < 1).
This compresses each component, forcing sampled actions to concentrate more tightly around each component's 's mean. However, merely shrinking the variances does not address the large-scale spread when mixture components are far apart. The inter-component variance remains, and noisy samples may still arise from small but non-negligible transitions among disparate components.
This phenomenon exacerbated when the choice the number of components is overspecified (see Figure aad further discussion in Appendix A. Furthermore, it has been shown by Carreira-Perpine [5] that GMMs in dimension d may have a mode that does not directly align with the mean of any component. In such situation, the down-scaling technique is likely to result in an insidious loss of the extra mode (see Figure
Mode-tracking via the Mean-Shift Algorithm We explored a principled way to reduce the variance of sampling from multimodal distributions by explicitly identifying the major modes and then sampling from those modes with an appropriate probability while controlling variance. An effective way to find the modes in GMM is via the mean-shift algorithm [5], which is an efficient fixed-point iteration algorithm (see Appendix B for further details). We propose multinomial distribution over the moded (m,) that respects the shape of the density:
Wj
Pr(m;) Eg Wy
(3)
where w = (mM)(m;) 1-11,1,
(4)
and H = V2 (ogPGMM((M))
(5)
Note that each H, admits a closed-form expression and can be computed Okk?² operations. One can then further inject noise to the sampled mode using a Gaussian noise with fixed variance or use the Hessian around the mode to account for the curvature of the distribution to determine the variance (see more details on this process in Appendix B). We emphasize that this mode-tracking algorithm can efficiently implemented on GPUs using common deep learning libraries and introduces minor overhead compared to the transformer forward pass.
4 Experiments
We evaluate the performance of Q-FAT through extensive experiments across 9 different tasks using five representative simulated robotics environments: PushT (7), Kitchen [TI) UR3 BlockPush,
'The number of modes J is hard to control theoretically [5].

---

ID GMM Variance Scaling
2D Learned GMM
2D Forcing of - 0.02
True
0.
GMM PDN
GMM modes
Componant
Component 2
1-
Dasky
R
0.2
-1-
-1
-10
Xg
Figure 2: The figures demonstrate the potential detrimental effects of down-scaling the component variances in a learned Gaussian Mixture Model (GMM). The left figure illustrates how down-scaling variance introduces irreducible variance proportional to the distance between component means, when the number of components (k) is misspecified (e.g., approximating unimodal distribution with two components). The two figures on the right show how reducing variance can lead to loss of multimodality, specifically causing the central mode to disappear in 2D two-mixture Gaussian that presents three modes.
BlockPush 10 and Multimodal Ant [26]. These environments cover the most common control signals in practice, namely position and velocity of robot's end effector and robot joint angles. A detailed description of each environment is provided in Appendix C We further ran preliminary experiments on an autonomous driving dataset (nuScene (4) and add the experimental results in Appendix B Our experiments are designed to address the following key research questions:
1. Is Q-FAT competitive with state-of-the-arr methods in both conditional and unconditional behavior imitation?
2. Can Q-FAT effectively capture the diversity of demonstration datasets while avoiding mode collapse?
3. What is the effect of different sampling techniques on the performance of Q-FAT?
4. How sensitive Q-FAT's performance the choice of the mixture components in the learned policy?
5. How does the inference speed of Q-FAT compare to VQ-BeT?
We find that Q-FAT achieves state-of-the-art performance across the majority of conditional and unconditional tasks, achieving state-of-the-ar balance between the quality and diversity of the generated actions. We report the results of Q-FAT using the variance reduced sampling (Section 3.2 with a factor of 10 and we discuss the implications of this choice in Section 4.3 Further training details are provided in Appendix D
summary of our results is presented in Table
4.1 Baselines and Performance Metrics
Baselines We compare Q-FAT against a range state-of-the-art baselines. We restrict our focus to comparing against diffusion and transformer-based policy learning methods. For unconditional tasks, the baselines include: (1) BeT [45), which performs action discretization using k-means clustering and models the action distribution with transformer decoder; (2) VQ-BeT [26]. extension of BeT that leverages vector quantization via a hierarchical Variational Autoencoder (VQ-VAE) [57] for action discretization; and (3) two variants of Diffusion Policy 77. namely a convolutional- implementation and transformer-basee the comparison include conditional versions of BeT and VQ-BET, well
implementation. For goal-conditional tasks, we extend additional baselines: BESO 42], denoising diffusion-based method with conditional variant (C-BESO) and classifier-free guided variant (CFG-BESO).
4.2 Quality and Diversity of Q-FAT's Actions
Evaluation Metrics We assess performance using task success rate and induced behavioral entropy. Task success is measured differently per environment: in Kitchen, Multimodal Ant and UR3 Block-
5

---

Table 1: Performance comparison between unconditional and conditional behaviors the task success rates. We report the performance of Q-FAT with variance down-scaling with a factor of 10
Unconditional Tasks
Environment
Metric
BeT
DiftPolicy-C
DiftPolicy-T
VQ-BeT
Q-FAT
PushT
0.39
0.73
0.74
0.78
0.80
Final IoU (-/1)
Image PushT
0.66
0.45
0.68
0.69
Kitchen Image Kitchen
3.07
2.62
3.44
3.66
3.84
Goals (+/4)
2.48
3.11
3.01
2.98
3.15
Multimodal Ant
2.73
3.12
2.90
3.22
3,30
UR3 BlockPush
1.59
1.83
1.82
1.84
1.99
Goals (-/2)
BlockPush
1.67
0.47
1.93
1.79
1.77
Conditional Tasks
Environment
Metric
C-BeT
C-BESO
CFG-BESO
VQ-BcT
Q-FAT
Cond. Kitchen
Goals (+/4)
0.15
3.75
3.47
3.78
3.78
Cond. UR3 BlockPush
Goals (-/2)
0.00
1.14
0.92
1.94
1.96
BeT
DiflesionPoliey-CC
Polleyt
VQ-BeT
Q-FAT (US)
423
0.90
2,66
4.00
3.58
3:35
3:12
2.55
2.0
0.00
Image Kitehco
BleckPesh
141
2.6-
1.99
1,59
p- E4-Emeapp
1.00
1.95
1.44
1.52
Figure 3: Overview of Q-FAT's behavioural entropy on unconditional behavior generation compared to baselines.
Push, the task success rate corresponds to the number of goals achieved per roll-out, while in PushT, task success is quantified using the final Intersection over Union (IoU) between the object's position and the target area. Behavioral diversity is measured by computing task completion sequence entropy (length for Kitchen and Multimodal Ant, length for UR3 BlockPush), reflecting variability in task completion. Baseline performance follows metrics from Lee et al. (26), and results are reported as the average success rate of the best checkpoint over 1000 episodes, including behavioral entropy where applicable.
Comparison Baselines Q-FAT achieves tate-of-thheart performance on success metrics, as shown in Table []] The performance gap observed with BeT can be attributed to the limitations of k-means clustering in high-dimensional spaces, which may hinder its ability to effectively capture complex action distributions. This is particularly exacerbated due to BeT attempting discretize the whole unconditional action space, and not the conditional one. In the case of VQ-BeT, the sequential nature of first predicting discrete action token followed by action correction step introduces two potential sources of error: inaccuracies token prediction and errors in the subsequent correction

---

Dataset
GMM Sampling
(c) o-scclld Sampling
(d) Mode Sampling
Figure 4: Sampling techniques from an 8-mixture Q-FAT policy multiroute environment 45 (a) Raw dataset with two pairs of equally likely paths (blue and red) from start target (green). (b)
Direct GMM sampling yields noisy samples due to the captured dataset noise. (c) Variance scaling (10 (10-& 8 rreduces per-component variance not the variance from inter-component distances. Mode sampling largely suppresses noise while preserving dataset multimodality.
mechanism. This places significant reliance on the vector quantization network, which is challenging optimize due the need to discretize the entire action space and the required two-stage training process - first for the quantization network and then for the policy. For diffusion-based policies, the absence of an explicit likelihood function makes it difficult control the variance of sampled actions, particularly in the presence of noisy data. This can be particularly problematic in fine ime-grained control tasks, where precise action execution is crucial for success. Consequently, while these approaches demonstrate strong performance in certain scenarios, they may struggle in tasks requiring high accuracy and reliability under noisy data. Q-FAT largely outperforms matches the behavioral entropy of all baselines (Figure 3), which assesses diversity across the generated motion.
4.3 Effects of Sampling Techniques
We assessed the effect of the different sampling techniques discussed in Section 3.2 down-scaling the component variances had a 7% and a 6.5% increase in the success performance metrics in the Kitchen environment and PushT, respectively, compared the GMM sampling, while only seeing marginal 2% drop in the diversity of the generated behaviors in the Kitchen environment.
We found that
The modes sampling matched the performance of the variance down-scaling the evaluation metrics. This could be attributed to the fact that the learned mixture components are sufficiently distant the environments explored, thus the effects discussed in Section 3.2 are not pronounced.
Despite quantitative metrics not showing significant differences between the variance down- scaling method and the mode sampling. we visualize the trajectories produced by both sampling techniques on a toy multi-route environment introduced by Shafiullah et al. [45] and see the effects of the irreducible variance in Figure 4 We further display a similar effect in the PushT environment in Figure 5
GMM Sampling
Variance Sampling
Sumpling
Figure 5: We visualize 400 trajectories generated from Q-FAT with 16 mixtures on the PushT environment with different sampling methods. One can see that mode sampling reduces the variance and does not produce trajectory artifacts.

---

Active Mixtures Evolution: UR3 BlockPush
Active Mixtures Evolution: Kitchen
Active Modes Distribation
0.9
0.9
UR3
Joorge Htais
No
0.7
0.7
0.5
0.5
0.2
0.3
0.3
0.0
20
40
80
100
40
60
80
100
Takking Progres O98
Traiskag Progres (S))
Figure 6: The left two plots display the evolution of the number of active mixtures during inference as training progresses, showing a clear pruning effect. On the rightmost plot, we display the distribution of the number of active modes at each timestep during inference for a converged model. One can see that trajectories remain largely unimodal.
4.4 Effect of the Number of Mixtures
We systematically evaluated the effect of varying the number of mixture components & in Q-FAT, where & controls the number of modes in the conditional action distribution for given state context.
Specifically, we considered k {2, 8, 16) in two benchmark domains, the Kitchen and UR3 BlockPush environments. Our results indicate that Q-FAT's success performance is highly robust across these values k; with negligible degradation in performance. observed an 18% increase in behavioral entropy when k was increased from to 4, after which entropy remained approximately constant as & continued to increase. In contrast, in the UR3 BlockPush environment, the behavioral entropy remained stable regardless of k. The key achieving this robustness is careful initialization.
the Kitchen environment, we
In particular, we found it important to initialize each mixture component's s mean to be distant from the others. Concretely, we placed these means on hyper-cube spanning the range [-1, 1] per dimension, maximizing their mutual separation, and assigned each component unit variance. The range [- 1] a reasonable choice, as actions are typically bounded in practice and normalized in our experiments.
To investigate how the mixture components were utilized during inference, we tracked the average number of active components over the course of training (see Figure 6 left). We observed that the model effectively prunes redundant mixture components as training progresses, with the pruning effect becoming more pronounced when more components are initially employed. Interestingly, in the Kitchen environment, which has a nine dimensional action space, less pruning occurs as increases compared to the UR3 environment, whose action space only two dimensional. A plausible explanation for this discrepancy lies in the dimensionality of the action space.
In higher dimensions, each Gaussian component covers proportionally smaller volume of the action space. Consequently, the model requires more components to adequately represent the conditional action distribution, and therefore prunes fewer mixture components. In lower-dimensional environments, by contrast, each component can cover a larger fraction the data, which makes aggressive pruning more viable.
We further plotted the distribution over the number of active modes for the Kitchen and the UR3 BlockPush environments. We observed that approximately 70% of the course of trajectory is uni-modal (Figure (6) right). This shows that multi-modalin is required only a minority of the time during rollout. This observation could be useful for future work to exploit this uni unimmodality computationally.
Inference Time Analysis
To compare computational efficiency, evaluated the inference times of Q-FAT and VQ-BeT on a GB MacBook Pro CPU, The analysis utilized models trained in the Kitchen environment, with hyperparameters for Q-FAT specified in Table 3 and for VQ-BeT described in (26). Q-FAT

---

is approximately 2x faster on the tested hardware due to its efficient action decoding head. The breakdown of inference time per model component is detailed Table 2
Table 2: Inference time breakdown for Q-FAT and VQ-BeT on the models used for the Kitchen
environment. Q-FAT has lighter action-decoding head, resulting in a signigicant speedup.
Component
VQ-BeT (ms)
Q-FAT (ms)
Shared Backbone (minGPT)
100
~00
Action Decoding Head
112
~011
Total Inference Time
~.2
1.1
The primary performance difference arises from the action decoding head. While both models use an identical backbone architecture, VQ-BeT's reliance on two separate 2-layer MLPs for action decoding creates a computational bottleneck. In contrast, Q-FAT's use single, efficient linear action decoding layer results in significant speedup.
5 Related Work
Offline Learning for Decision Making Learning to act from offline data has received significant attention over the years due to its vast potential in practical applications. This line of research
can be divided into two categories: offline reinforcement learning [25 27 22 6 and imitation
learning [34 37 38 [3]. In offline reinforcement learning, on the one hand, the goal is learn policy from demonstrations with suboptimal behaviors, which necessitates the access to rewards in the demonstration dataset. Imitation learning, on the other hand, assumes that the trajectories are collected from near optimal behaviors, avoiding access to reward signals in the collected datasets. This increases the attractiveness of imitation learning practical setups, since reward signals are hard to define [9, 23 Q-FAT falls under the framework of behavioral cloning, which is a form of imitation learning where the policy learning is treated as supervised learning problem.
Generative Models for Cloning Behavior Complex generative models have been employed to capture the full data diversity due to the complex and multimodal behavior in human and robot datasets. Earlier work has explored the use of Gaussian processes 46]. energy-based models
10 and generative adversarial networks 21 as policy parametrizations. More recently, more
and 36). Unlike other transformer-based policies, Q-FAT uses a continuous parametrization of the output by predicting GMM. Note that this idea has been explored before using LSTM backbones [29 45), however, it has been shown to produce poor performance [45] 26).
attention has been directed towards more powerful models, such as transformers [55]
26 32, 55
diffusion models 154 40
which we attribute to the instability of training LSTM networks and their limited capacity to model long sequences
Discussion and Future Work
In this work, introduced Q-FAT, quantization-free action transformer that achieves state-of-the- art performance across various simulated robotics environments, opening new research directions in both behavioral modeling and reinforcement learning. Compared prior work (26), Q-FAT overcomes the challenging non-differentiable step in action discretization while maintaining or improving performance.
By integrating Q-FAT into end-to-end policy learning pipelines, future research could evaluate its influence sample efficiency and final task performance. Another promising area involves incor- porating Bayesian priors into Q-FAT's action distribution estimates, which could facilitate active exploration in complex, high-dimensional settings. For fine manipulation tasks, researchers could further investigate coarse-to-fi sampling strategies based on Gaussian mixture model represen- tations, potentially improving both exploration breadth and control. Finally, extending Q-FAT to non-Euclidean action spaces, such as those with Riemannian geometry, may enable more accurate and natural representations for tasks like legged locomotion dexterous manipulation.

---

7 Limitations
While we have demonstrated that Q-FAT is effective in learning complex action distributions, the underlying GMM loss function used to train the policy assumes a Euclidean action space. In some environments (e.g., humanoid robots), the action space has a Riemannian geometry, which could potentially result in difficulty during learning. This can be counteracted by mapping actions into latent Euclidean space using an action autoencoder [19] and perform the GMM loss in the latent space. However, unlike previous work that uses discrete action encoders [26 45 Q-FAT allows the flexibility of using continuous latent spaces, allowing end-to-end differentiability of the joint encoder and policy models. While this is exciting research direction, leave this for future work.
8 Impact Statement
Our work tackles simulated continuous control tasks whose downstream consequences are important for robotics. The probabilistic nature of our method may impact the safety of the generated behavior and should be further studied and tested. However, unlike prior work, our approach enables uncertainty estimates directly in the action space, which have the potential to improve safety by providing uncertainty quantification.
Acknowledgment
We would like to thank Youssef Saied for the valuable discussions throughout the project and Nicolas Nguyen for reviewing the manuscript. C. V. and Z.S. are funded by the German Research Foundation (DFG) under both the project 46880671 14 of the Emmy Noether Programme and under Germany's Excellence Strategy EXC number 2064/1 Project number 390727645. M.M. is funded by the German Research Foundation (DFG) under the project 456587626 of the Emmy Noether Programme. We also thank the international Max Planck Research School for Intelligent Systems (IMPRS-IS) for the support.
References
[1] Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum posteriori policy optimisation. arXiv preprint arXiv: arXiv:/806.000 1806.06920, 2018.
[2] Karl J Åström. PID controllers: theory, design, and tuning. The International Society of Measurement and Control, 1995.
[3] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33: 33:1877-1901, 2020.
[4] Holger Caesar, Varun Bankiti, Alex Lang, Sourabh Vora, Venice Erin Liong, Qiang Xu, Anush Krishnan, Yu Pan, Giancarlo Baldan, and Oscar Beijbom. nuscenes: A multimodal dataset for autonomous driving. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 11621-11631, 2020.
[5] Miguel A. Carreira-Perpinan Mode-finding for mixtures of Gaussian distributions. IEEE Transactions on Pattern Analysis and Machine Intelligence, 22(11):1188-133, 2000.
[6] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Picter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in Neural Information Processing Systems, 34:15084-15097, 2021.
[7] Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, 2023.
10

---

[8] Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Advances in Neural Information Processing Systems, 32, 2019.
[9] Rati Devidze, Goran Radanovic, Parameswaran Kamalaruban, and Adish Singla. Explicable reward design for reinforcement learning agents. Advances in Neural Information Processing Systems, 34:20118 11220118-201111 2021.
[10] Pete Florence, Corey Lynch, Andy Zeng, Oscar Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168, 2022.
[11] Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving Long-horizon tasks via imitation and reinforcement learning. arXiv 2019.
preprint
[12] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for im- age recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770-778, 2016.
[13] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning, Advances in Neural Information Processing Systems, 29, 2016.
[14] S Hochreiter. Long short-term memory. Neural Computation, MIT-Press, 1997.
[15] Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley, Stochastic variational inference. Journal of Machine Learning Research, 2013.
[16] Matthew W Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Mom- chev, Danila Sinopalnikov, Piotr Stanczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, et al. Acme: A research framework for distributed reinforcement learning. arXiv preprint arXiv:2006.0099 2020.
[17] Jiaheng Hu, Rose Hendrix, Ali Farhadi, Aniruddha Kembhavi, Roberto Martín-Marthn Peter Stone, Kuo-Hao Zeng, and Kiana Ehsani. Flare: Achieving masterful and adaptive robot policies with large-scale reinforcement learning fine-tuning. arXiv preprint arXiv:2409./6571 2024.
[18) Minyoung Huh, Brian Cheung, Pulkit Agrawal, and Phillip Isola. Straightening out the straight- through estimator: Overcoming optimization challenges in vector quantized networks. In International Conference on Machine Learning, pages 14096-14113, 2023.
[19] Diederik Kingma, Max Welling, al. Auto-encoding variational bayes, 2013.
[20] Alexander Kolesnikov, André Susano Pinto, Lucas Beyer, Xiaohua Zhai, Jeremiah Harmsen, and Neil Houlsby. Uvim: A unified modeling approach for vision with learned guiding codes. Advances in Neural Information Processing Systems, 35:26295- 26308, 2022.
[21] Alex Kuefler, Jeremy Morton, Tim Wheeler, and Mykel Kochenderfer. Imitating driver behavior with generative adversarial networks. In IEEE Intelligent Vehicles Symposium, pages 204 114,211, 2017.
[22] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33: 1179-1191, 2020.
[23] Aviral Kumar, Joey Hong, Anikait Singh, and Sergey Levine. When should we prefer offline reinforcement learning over behavioral cloning? arXiv preprint arXiv:2204.056 2022.
[24] Adrian Lańcucki. Jan Chorowski, Guillaume Sanchez, Ricard Marxer, Nanxin Chen, Hans JGA Dolfing, Sameer Khurana, Tanel Alumäe, and Antoine Laurent. Robust training of vector quantized bottleneck models. In International Joint Conference on Neural Nenworks, pages 1-7, 2020.
[25] Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State-of-the-art, pages 45-73. Springer, 2012.

---

[26] Seungjae Lee, Yibin Wang, Haritheja Etukuru, H Jin Kim, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. Behavior generation with latent actions. arXiv preprint arXiv:2403.03181 2024.
[27] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 2020.
[28] Hao Ma, Dieter Büchler, Bernhard Schölkopf, and Michael Muehlebach. Reinforcement learning with model-based feedforward inputs for robotic table tennis. Autonomous Robots, 47 (8):1387-1403, 2023.
[29] Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkarni, Li Fei-Fei, Silvio Savarese, Yuke Zhu, and Roberto Martín-Martnn What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2/08.03298, 2021.
[30] Jiageng Mao, Yuxi Qian, Junjie Ye, Hang Zhao, and Yue Wang. Gpt-driver: Learning to drive with gpt. arXiv preprint
2023.
[31] Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite scalar quantization: VQ- made simple. In International Conference on Learning Representations, 2024.
[32] Atharva Mete, Haotian Xue, Albert Wilcox, Yongxin Chen, and Animesh Garg. Quest: Self- supervised skill abstractions for learning continuous control. arXiv preprint arXiv:2407. 15840, 2024.
[33] Michael Muchlebach and Raffaello D' Andrea. Nonlinear analysis and control of a reaction- wheel-based 3-d inverted pendulum. IEEE Transactions on Control Systems Technology, 25(1): 235-246, 2017.
[34] Jyothish Pari, Nur Muhammad Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation. arXiv preprint arXiv:2112.00511, 2021.
[35] Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsur- prising effectiveness of pre-trained vision models for control. In International Conference on Machine Learning, pages 17359-17371, 2022.
[36] Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et al. Imitating human behaviour with diffusion models. arXiv preprint arXiv:2301. 2023.
[37] Xue Bin Peng, Pieter Abbeel, Sergey Levine, and Michiel Van de Panne. Deepmimic: Example- guided deep reinforcement learning of physics-baser character skills. ACM Transactions On Graphics, 37(4): 2018
[38] Xue Bin Peng, Ze Ma, Picter Abbeel, Sergey Levine, and Angjoo Kanazawa. Amp: Adversarial motion priors for stylized physics- hased character control. ACM Transactions on Graphics, 40 (4):1-20, 2021.
[39] Jan Peters and Stefan Schaal. Reinforcement learning of motor skills with policy gradients. Neural Nenworks, 21(4):682-697, 2008
[40] Michael Psenka, Alejandro Escontrela, Pieter Abbeel, and Yi Ma. Learning diffusion model policy from rewards via q-score matching. arXiv preprint arXiv:2312. 2023.
[41] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training OpenAl, 2018. URL httpttt:/ccc..opeeeee com/ Research-covers/laagg http:://aaaggggguuuuuuuurrreaald aawgggaaagggguddrrtttddiiinnrrpp pdf
[42] Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304.02532, 2023.
12

---

[43] Stefan Schaal. Learning from demonstration. Advances in Neural Information Processing Systems, 9, 1996.
[44] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint 15K.::18000002438,, 2015.
[45] Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with stone. Advances in Neural Information Processing Systems, 35:22955-22968 2022.
[46] A.P. Shon, K. Grochow, and R.P.N. Rao. Robotic imitation from human motion capture using Gaussian processes. In JEEE-RAS International Conference on Humanoid Robots, pages 129-134, 2005.
[47] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13450 2020.
[48] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01446, 2023.
[49] Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in Neural Information Processing Systems, 12, 1999.
[50] Michael Tschannen, Cian Eastwood, and Fabian Mentzer. Givt: Generative infinite-vocabulat vocabulary transformers. In European Conference on Computer Vision, pages 292-309. Springer, 2024.
[51] Michael Tschannen, André Susano Pinto, and Alexander Kolesnikov. Jetformer: An autore- gressive generative model of raw images and text. In International Conference on Learning Representations, 2025.
[52] Aaron Van Den Oord and Oriol Vinyals. Neural discrete representation learning, Advances in Neural Information Processing Systems, 30, 2017.
[53] A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 30,
2017.
[54] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:228..06193, 2022.
[55] Hongtao Wu, Ya Jing, Chilam Cheang, Guangzeng Chen, Jiafeng Xu, Xinghang Li, Minghuan Liu, Hang Li, and Tao Kong. Unleashing large-scale video generative pre-training for visual robot manipulation. arXiv preprint 132.13231,.13139, 2023.
[56] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXi::/911.11330 2019.
[57] Neil Zeghidour, Alejandro Luebs, Ahmed Omran, Jan Skoglund, and Marco Tagliasacchi. Soundstream: An end-to-end neural audio codec. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 30:495-507, 2021.
[58] Yuke Zhu, Ziyu Wang, Josh Merel, Andrei Rusu, Tom Frez, Serkan Cabi, Saran Tunyasuvu- nakool, János Kramár, Raia Hadsell, Nando de Freitas, et al. Reinforcement and imitation learning for diverse visuomotor skills. arXiv preprint arXiv: 1802.8000266 2018.
[59] Jasan Zughaibi, Bradley Nelson, and Michael Muchlebach. Dynamic electromagnetic naviga- tion. Robotics and Automation Letters. 10(6):6095.. 0000600056066....0000000 2025.
13

---

A GMM Properties
GMM Moments
Let X be random variable from a k-mixture GMM distribution in R The probability density function can then be written as:
p(x) = Σ TN p4 Ei),
(6)
where N(x M4g E,) is the multivariate Gaussian distribution for the i-th component. The mean and the covariance of X are given by:
E
Tip4;
1=1
Cov
Σ
Ti E
(14
E ( B[X))),
(8)
where are the mixing coefficients, and μ. Σ; are the mean and covariance parameters for the i-th Gaussian component. Note that while the mean of the GMM is a weighted sum of the component means, the covariance is not. The total covariance reflects both the inherent variability within each component and the variability introduced by the differences component means.
B
GMM Mode Finding and Sampling
Mode Finding
To find the modes of GMM, we compute the gradient of the probability density function with respect to and equate it to zero:
Vp(x) Σ n. VN(x p., E.) =
(9)
where the gradient of the multivariate Gaussian distribution is given by
VA((X u. E.) [x pai Σ.)E, (m, x).
(10)
We obtain therefore:
TN 14 Σ,)Σ, (H:
0.
(11)
We identify the modes of the GMM by computing all critical points x" that satisfy (1). Moreover, ensure that a critical point x* is indeed a mode, the Hessian of the density function at that point must be negative definite. The Hessian the GMM density is given by
k V p((()) Σ m. V>N(x M4: E.),
(12)
where the Hessian of the multivariate Gaussian distribution is:
V p4 Ex) [x μ, Σx) Σ (H, x)(,, - x) Tx, -----1]
(13)
Therefore, the Hessian of the GMM density becomes
p(x()) Σ
Tx p4: E.) Σ, (1: x)(4; xx) Σ -
(14)
and critical point x" mode if: Vp(x*) 0 and V²p(x") is negative definite.
14

---

While GMMs are often parameterized by finite number of components, the resulting distribution can exhibit a richer structure. particularly in digees-iimesiioall spaces. Specifically, in dimensions greater than two (d >> 2). GMM can possess more modes (local maxima) than the number of its constituent components [ This phenomenon arises due to the interplay between the covariance structures and the relative positions of the Gaussian components. In higher dimensions, the overlapping regions of different components can create multiple peaks the probability density function that cannot be attributed to individual components.
Mean-Shift
We start by solving for x as follows:
K
F.
Nxx
Mic Ex) DT
Σ
Nx
M4, Eg)
p4
(15)
To obtain the mean-shift update in practice, one can treat the solution to More precisely, we define the operator T 12d + Rd by
15 as fixed-point iteration.
k
T(x)
Σ 1
N u. E.) 2T
Σ
N H.: E. Ef P..
(16)
Then, starting from initial guess x ("), the mean-shift procedure updates the estimate of via
(f+1) = T(x
The fixed-point iteration has the following interpretation: each iteration shifts the current iterate after t iterations closer to the modes -tishy regions) of the underlying density defined by the mixture Gaussians. In practice, it has been shown that initializing x() with the component means μ, results in fast convergence [5]. which also observe in experiments. To account for the fact that the number of modes could exceed the number mixture components, we add extra initialization points from the convex-hull of the component centroids capturing the most important modes.
Mean-Shift with Diagonal-Covariane Components
An important special case arises when the covariance matrices E; are diagonal. In that scenario, let
diag((o2,,. L.a.)
Since each D; is also diagonal (with entries 1/02., along the diagonal), the vector and matrix oper- ations inside T(x) decouple across dimensions. Hence the inverse of Σ N> H4 Σ,) E,- reduces to a diagonal matrix whose components can be computed efficiently. As result, each coordinate of the updated point can be determined independently. making the meaf- iteration particularly fast on modern hardware.
Modes Sampling
To sample the modes proportional to their coverage of the data support, we approximate the integral of the GMM density the local neighborhood of each mode using Laplace approximation and use the value of the integral around each mode to compute the relative weight of the mode. We summarize the full mode-findin finding and sampling procedure Algorithm
C Environment Details
In our experiments, we evaluated Q-FAT in three distinct environments: Kitchen, PushT, and UR3 BlockPush. the following, we provide brief description of each environment and its variants.
Franka Kitchen: This environment involves a Franka Panda robotic arm with a nine dimensional action space, designed for multi-task manipulation in a simulated kitchen
15

---

Algorithm Mode Extraction and Sampling
Inputs: GMM parameters { p4, Et tolerance e; rejection threshold 0; max iterations
max it; number of initializations n init;
2: Initialize Mode Set: MM
3: Initialize Seed Set:
4: X'
Uniform Uniforn_Sampaaaeee (Cmnpee(Convex_Hullll n_iit))
5: for E X' do
6: Iteration counter i 4-
7: Define the operator (-) as in Equation 16
8: repeat
9: Update via fixed-point iteration: + T(x)
10:
Increment
11:
until i max it or ||x Xotall <
12:
Compute H V px)
13:
if max (exxengenvvalus (H)) < 0 then
14:
Merge nearby modes within a distance threshold, keeping the highest-probabiliy mode Add mode to M
15:
16: end if
end for
18:
19: Compute Mode Weights:
log ((x) 1/2 O...../M 20: wy = p(xx)) VV²
21: Sample mode: Sample j Categorical ((w;)]=0)
22: Return: Xj
The dataset consists 566 human-collected demonstrations, where each trajectory completes a subset of four out of seven possible tasks in varying orders. For the image-based version of the environment, we rendered the collected trajectories into 112x112 images.
setting [1].
PushT: The PushTT environment focuses pushing, T-shaped block to target position on a table using two dimensional end-effector velocity control 7 The dataset includes 206 human-collected demonstrations.
UR3 BlockPush: In this task, a UR3 robotic arm is tasked with moving two blocks to designated goal circles on table 10. The dataset exhibits multimodality, as the blocks can be moved in either order. In the unconditional setting, we evaluate whether both blocks reach their respective goals. In the conditional setting, the model is provided with the target positions of the blocks, and performance is assessed based on the order in which the blocks reach their goals.
BlockPush: The goal of this task is moving two (red and blue) blocks two targets [T]] The blocks can be moved in either order and can be put into either of the targets. While this environment is similar to UR3 BlockPush, it exhibits more stochasticity due the random initialization of the block positions, and the targets being randomly translated and rotated. The dataset contained 1000 demonstrations.
Multimodal Ant: The Multimodal Ant environment
26 is modification of the Ant
environment introduced by Schulman et al. 44 where the goal of the ant is to visit four distinct locations placed vertices of square. The multimodality in the dataset arises from the different order the goals are visited. The dataset contains 600 human-collecte trajectories
16

---

D
Training Details
For our experiments, used the minGPT backbone the decoder-only transformer implementa- tion. The hyperparameters used for each of the environments are detailed in Table 3
D.1 Architectural Design Choices
The policy was trained using teacher forcing with causal attention mask, this provided stronger learning signal compared to computing the loss only on the last action given a context states. We found that using teacher forcing and employing full attention mask significantly degraded performance. We further experimented with feeding back the sampled actions into the model's 's input by concatenating states and actions into single feature vector. This approach degraded policy performance in simulated environments, likely due to the model overfitting to the highly correlated previous actions in the context while ignoring state information. Injecting Gaussian noise into the previous actions was necessary to achieve reasonable performance in this setting; however, this provided no discernibee benefit compared to policy without action feedback in our experiments.
D.2 History Masking
During training, we observed instances where both validation and training losses decreased, yet the environment rewards declined. We attribute this phenomenon causal confusion [8], where the model overfits the recent sequence of states. Instead of learning true causal relationships, the model relies temporally correlated but non-causal patterns predict future actions, leading to suboptimal performance. To mitigate this, we introduced history masking, randomly masking the context (excluding the current state) with certain probability during training. Applying this technique in the Kitchen, Multimodal Ant, and BlockPush environments (with masking probabilities of 0.7, 0.3, and 0.3, respectively) led to improved correlation between environment rewards and validation loss.
D.3 Image-based Environments
For Image PushT and Image Kitchen, fine-tuned ResNet18 encoder [12], extracting feature maps from the first four layers to construct low-level features (35]. These features were then spatially average-pooled and fed into the transformer with dimensionality of 1024. For Image PushT, applied data augmentation (color jitter, random cropping, and random gray-scaling with probability 0.5), which proved crucial for preserving the invariance of the pre-trained ResNet and preventing overfitting to our dataset. Conversely, for Image Kitchen, we found data augmentation deteriorated performance.
D.4 Goal Conditioning
In goal-conditional tasks, we conditioned the model on a sequence of future states equal in length to the state history. During inference, we randomly sampled reference trajectory from the dataset. The agent's reward was computed only if it achieved the goal sequence in the same order as the reference trajectory.
D.5 Evaluation
All training datasets were normalized using min max scaling to ensure state and action features lie within the range [-1, 1, 1]. Following Lee et al. (26). Shafiullah et al. [45], split the data into 95% for training and 5% for validation. The policy was evaluated periodically during training using a small number of environment roll-outs (20 to 50). The models achieving the lowest validation loss and highest success metrics were selected for reporting.
D.6 Model Size
For fair comparison with VQ-BeT (26.. we ensured our models have comparable (or smaller) number of trainable parameters, on the order of 1 encoder hyperparameters to account for the additional capacity utilized by their hierarchical vector
106. We slightly deviated from their transformer
17

---

Hyperparameter
P
Image PushTT'
Kitchen
UR3
Multimodal Ant
BlockPush
nuScenes
Layers
6
6
6
6
6
4
6
Attention heads
8
8
8
8
8
4
8
Embedding dimension
128
128
128
128
128
72
128
Dropout probability
0.1
0.1
0.1
0.1
State history size
5
5
10
10
10
3
Action horizon
5
5
1
6
Training epochs
1000
1000
1200
400
1000
1200
1000
Batch size
256
256
128
2048
128
256
128
Number mixtures
2
Maximum learning rate
le-3
le-3
le-3
le-4
le-3
lc-4
1c-4
Minimum learning rate
lc-6
lc-6
le-6
1c-7
Learning Rate Schedule Optimizer
Cusine
Cusine
Cosime
Cosine
Adam
Adam
Adam
Adam
Adam
Adam
Adam
Table 3:
Environment hyperparameters.
quantization autoencoder for action discretization. Training hyperparameters, such batch size and learning rates, were adjusted to accommodate the differences in loss functions between the methods. Further details are provided in Table and Table in Lee et al. [26).
D.7 Hardware
Experiments were conducted on a heterogeneous cluster, making precise hardware control infcasible. However, all experiments Training a single model typically took 4-8 hours, depending on dataset size and the frequency of environment evaluations for validation.
single desktop-grade GPU with most 32 GB of memory.
E
nuScenes Experiment
To evaluate our model's applicability beyond robotic manipulation, we use the nuScenes dataset, large-scale benchmark for autonomous driving. We utilize the object-centric, preprocessed version of the dataset from Mao et al. [30], and follow the exact preprocessing and tokenization steps for driving data as detailed in the VQ-BeT paper [26). The model's input a of tokens representing the driving mission (e.g., turn left), the ego-vehicle's current state (velocity, acceleration). its recent trajectory history, and the state surrounding objects (position, class). The task is to predict the ego-vehicle's trajectory over the next six timesteps. Success is evaluated using two primary metrics: the average L2 distance between the predicted and ground- toound-rruth trajectories to measure accuracy, and the collision rate to assess the safety of the generated path.
For this sequential prediction task, found necessary adapt model's architecture. A naive approach predicting a single high-dimensional vector concatenating all future waypoints resulted in degraded performance. This is likely because the model is forced learn direct mapping highly complex and multimodal joint distribution of future states. Such a method struggles to enforce temporal consistency, often producing kinematically implausible trajectories.
Table 4: Preliminary results from our initial experiments on the nuScenes autonomous driving benchmark. Lower values are better.
Method
L.2 Error (m)
Collision Rate (%)
Diffusion
0.96
0.44
VQ-BeT
0.73
0.29
Q-FAT (ours)
0.75
0.37
Consequently, we adopted an autoregressive prediction scheme, where the model forecasts one way- point at time, conditioned on its previous predictions. This approach simplifies the learning problem by factorizing the joint distribution into a sequence of more manageable conditional distributions. By doing so, the model implicitly learns the transition dynamics of the environment, ensuring that the generated trajectory temporally coherent and physically plausible. Furthermore, instead of predicting the absolute coordinates of each waypoint, our model predicts the delta, or displacement,
18

---

from the previous waypoint. This encourages the model to learn policy that is invariant to the absolute frame of reference, which has been shown improve generalization.
The results from these initial experiments are presented Table 4 These findings show that Q-FAT achieves performance competitive with the VQ-BeTT baseline, demonstrating that method is effective in this challenging domain. The training hyperparameters can be found in Table
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We introduce Q-FAT and evaluate it extensively on common bench marks, demonstrating SOTA performance.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: See Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification:
We do not present any theoretical claims.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [Yes]
Justification:
Our code and datasets will be made publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/t.pss.ccppullc/guuii / SwweleeeeooSSSSeeeeeeeeeee
www.SSubbisiooooiiiooooooooo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: See Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification; It was computationally expensive train multiple models for each environment to assess the modeling uncertainty. As for the aleatoric uncertainty in the environments, it was negligible since we ran evaluations on 1000 episodes and averaged the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
22

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: We will open source our code with public licenses. We further clearly stated the societal impact in Section g
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Section g
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairesss considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
23

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We train specialized models for simulated robotics tasks, which does not require safeguards.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We ensured that all the data and code that was used was made publicly available for research use.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi theode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
24

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
Justification:
We will open source our code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our work does not include human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does include human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
25

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: We did not use LLMs to conduct this research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy for what should or should not be described.
26


Reviews:
[{"review_id": "wx9nznmae4", "reviewer": "Reviewer_j3Wb", "summary": "This paper introduces Quantization-Free Autoregressive Action Transformer (Q-FAT), a novel policy representation for continuous-control imitation learning that avoids the common step of discretizing actions. Instead of tokenizing actions via VQ-VAE or k-means clustering, Q-FAT directly predicts the parameters of a Gaussian mixture model (GMM) at each timestep with a decoder-only transformer backbone. The policy outputs mixture weights, means, and (diagonal) variances, enabling exact likelihood evaluation, autoregressive sampling, and principled variance-reduction techniques (variance scaling and mode-tracking via mean-shift). Empirically, Q-FAT matches or exceeds the state of the art on nine simulated robotics tasks (e.g., PushT, Kitchen, UR3 BlockPush, Multimodal Ant), improving both success rates and behavioral diversity while maintaining efficient inference. The authors also analyze robustness to the number of mixtures, illustrate pruning dynamics, and discuss practical training details (history masking, ResNet encoder for vision, initialization).", "questions": "1. Could you isolate the benefit of continuous GMM outputs by comparing a discrete-token transformer with identical depth/width to Q-FAT?\n\n2. How does Q-FAT perform on real-robot or higher-dimensional (e.g., 17-DoF humanoid) settings? Can the GMM scale gracefully as action dimensionality grows?\n\n3. What is the wall-clock overhead of mean-shift sampling in practice? A small timing study (e.g., ms per step) would help assess real-time applicability.\n\n4. How sensitive are the sampling techniques to varying noise levels in demonstrations? An ablation where dataset noise is artificially increased could inform robustness.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Quality\n\n**Strength**: solid experimental validation across diverse benchmarks, with thorough comparisons to transformer- and diffusion-based baselines (BeT, VQ-BeT, DiffusionPolicy, BESO); state-of-the-art task success rates and entropy metrics demonstrate both accuracy and multimodality.\n\n**Weakness**: limited ablation on the impact of transformer depth/width; hard to disentangle gains from architecture vs. continuous parametrization; no real-robot or higher-dimensional continuous tasks (e.g., humanoid control) to test scalability beyond simulation.\n\n# Clarity\n\n**Strength**: well-structured presentation: clear problem motivation, method description, and algorithmic details (Equations 1-5, Algorithm 1); helpful figures (GMM sampling artifacts, mode pruning evolution).\n\n**Weakness**: appendices are dense: readers may struggle to find key hyperparameters or environment details without deeply navigating Appendix D; limited discussion connecting latent continuous policies to recent diffusion-based approaches beyond performance tables.\n\n# Significance\n\n**Strength**: eliminates a major pipeline complexity (action quantization), simplifying training and deployment; enables exact likelihoods and fast, single-pass inference, addressing diffusion-policy latency issues.\n\n**Weakness**: empirical gains, while consistent, are incremental (e.g., +2-7 % success), raising questions about broad impact across more challenging domains; the practical value of variance-reduction techniques might be context-specific (noisy vs. clean data).\n\n# Originality\n\n**Strength**: novel adaptation of Generative Infinite-Vocabulary Transformers (GIVT) to continuous control, directly modeling GMM outputs; introduction of mean-shift mode sampling in action transformers.", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thorough evaluation and insightful feedback on our work. We will update the appendix in the updated manuscript to make it less dense as you suggested. We further answer your questions below:\n\n- **Isolating the benefit of continuous GMM outputs by comparing a discrete-token transformer with identical depth/width**: We have indeed isolated the benefit of the continuous GMM parameterization by ensuring that we have the same or less trainable parameters than the discrete-token transformer baseline (VQ-BeT). Having an identical width and depth would give our model a disadvantage, since it does not have a separate quantization network as in VQ-BeT.\n\n- **Performance beyond simulated tasks**: We have conducted new experiments on the nuScene [Caesar et. al, 2020] autonomous driving benchmark as suggested by reviewer DqgJ. Our preliminary results show that Q-FAT achieves performance competitive with the state-of-the-art VQ-BeT baseline, demonstrating that our method is effective in complex, real-world-adjacent scenarios. Here is the table with preliminary results from our initial experiments (lower is better):\n| **Method** | **L2 Error (m)** | **Collision Rate (%)** |\n| :----------- | :--------------- | :--------------------- |\n| Diffusion | 0.96 | 0.44 |\n| VQ-BeT | **0.73** | **0.29** |\n| Q-FAT (ours) | 0.75 | 0.37 |\n\n We plan to add this experiment to our revised manuscript.\n\n- **GMM scaling to higher dimensions**: The GMM head's parameter count scales linearly with the action dimension when using a standard diagonal covariance matrix, meaning it does not hinder the model's application to higher-dimensional problems. Furthermore, recent work [Tschannen et. al 2024] has demonstrated scaling GMMs to high dimensional image modeling with an autoregressive transformer (in a 128-dimensional latent space with 1k mixture components), demonstrating high sample fidelity. This indicates that scaling GMMs to complex distributions in high-dimensional spaces, such as action spaces, is straightforward.\n\n- **Inference Time Analysis**: We have performed an inference time comparison between the sampling algorithms on a trained model from the Kitchen environment, where the model has 4 mixture components. Here is the table of the results:\n| **Method** | **Time (milliseconds)** |\n| :------------ | :---------------------- |\n| GMM Sampling | 1.061 ± 0.125 |\n| Mode Sampling | 1.550 ± 0.314 |\n\n The performance was tested on a 16 GB Macbook Pro CPU. We will add a more comprehensive study in the revised manuscript across the different environments.\n\n- **Sensitivity of the sampling techniques to varying noise levels in demonstrations**: Figures 4 and 5 illustrate the impact of various sampling techniques. In low-noise environments, all techniques exhibit comparable performance. However, with standard GMM sampling, an increase in noise tends to significantly destabilize the generated paths. For variance-reduction and modes sampling, the quality of actions is highly contingent on the underlying structure of the ground truth distribution and the number of mixtures used to model it, which we visually represented in Figure 2.\n\n--------------------------------------\n- [Caesar et. al, 2020] nuScenes: A Multimodal Dataset for Autonomous Driving\n- [Tschannen et. al 2024] JetFormer: An Autoregressive Generative Model of Raw Images and Text"}, {"review_id": "VZsxypP62p", "reviewer": "Reviewer_vHHX", "summary": "In this work, Q-FAT is introduced as an alternative to action models that are either diffusion based (thus slow), and autoregressive-based (thus relying on vector quantization). The approach is based on a recent Generative Infinite-Vocabulary Transformer (GVIT), which allows for continuous parametrization through a mixture of gaussians. This allows for quick inference and for sampling to happen, which is essential when the outputs of the model must be used as action policy. In the work, two sampling techniques are explored to reduce jitteriness in the predicted trajectory, showing that both are successful. The experiments are on 9 different tasks, where the method display state-of-the-art performance.", "questions": "How would you address the scaling question? This model is transformer-based, however, it relies on GMM. Will this hinder the scaling performance?\n\nPlease alter the plots to have larger labels, and improve the clarity of Figure 1.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Quality**. The methodology is well-introduced and justified, by explaining the downsides of previous methods. The issues with sampling strategies are exemplified, and later explored in the experiments. The model is evaluated on diverse tasks, showing good performance against relevant competing approaches.\n\nGiven the similar performance of the two sampling techniques it is not clear to me why they are both introduced. What are the true advantages of one over the other? Could this be demonstrated, or at least explained more clearly?\n\nThe issue of the scalability of such approach is also important and is not addressed.\n\n**Clarity**. The manuscript is clear and well written. The metrics are well-explained and it is clear to see when the method outperforms the others, and also by how much. Model behavior is exemplified with plots.\n\nThe plots should have larger labels, that are consistent with each other. \nThe first figure could be more clear, the exact environments are not clear, and the masking used in the transformer is also not specified. Finally, for Figure 1, it would be important to also highlight that special sampling techniques are used to get smooth trajectories. \n\n**Significance**. This is a significant contribution to the field. The model proposed addresses specific issues, and could be used as inspiration for further development, for example when combined with LLMs.\n\n**Originality**. The proposed approach takes a previously-proposed model, and adapts it in an original way for robotic applications. This is an original piece of work.", "was_revised": true, "final_justification": "This work has shown a novel and successful application of GVIT to robotics, which will be useful to the community. \nThe methods are mostly well described and the main concerns were addressed in the rebuttal. The results are positive and support the claims. This is why I believe this paper should be accepted.", "rebuttal": "Thanks for your thorough evaluation and insightful feedback on our work. We will update the description of Figure 1 as you suggested. We further answer your questions below:\n\n- **Motivation for Two Sampling Techniques**: The two techniques were introduced to address trajectory 'jitter' that we observed in preliminary experiments, even after significantly downscaling the GMM variance. This jitter arose from the model rapidly switching between different predictive modes between timesteps. Mean-shift sampling enforces temporal consistency and produces smoother trajectories, as shown qualitatively in Figures 4 and 5. We will clarify this motivation in the revised text.\n\n- **Scalability**: In terms of model and data scalability, autoregressive transformers are the most scalable model known to date, as demonstrated by frontier level LLMs. As for the GMM, the GMM head's parameter count scales linearly with the action dimension when using a standard diagonal covariance matrix, meaning it does not hinder the model's application to higher-dimensional problems. Furthermore, recent work [Tschannen et. al 2024] has demonstrated scaling GMMs to high dimensional image modeling with an autoregressive transformer (in a 128-dimensional latent space with 1k mixture components), demonstrating high sample fidelity. This indicates that scaling GMMs to complex distributions in high-dimensional spaces, such as action spaces, is straightforward.\n\n- **Masking**: We use causal masking as described in appendix D. We will highlight this more in our revised manuscript.\n\n---\n### References\n[Tschannen et. al 2024] JetFormer: An Autoregressive Generative Model of Raw Images and Text"}, {"review_id": "dvQ95vnEFB", "reviewer": "Reviewer_DqgJ", "summary": "This paper adopts the GIVT model to represent the control policy of robotic systems. By using a Gaussian Mixture Model (GMM) as the output distribution, GIVT avoids the need for discretizing actions. The authors also propose a mode-tracking sampling technique to reduce variance. Experiments across five robotic environments demonstrate that Q-FAT achieves state-of-the-art performance.", "questions": "•\tReal-world robotics and autonomous driving experiments are missing. While real robotic experiments might be difficult to conduct at this stage, it is important to include at least some experiments in autonomous driving scenarios (e.g., CARLA, nuScenes) to further demonstrate the practical utility of the proposed method.\n•\tThe mode-finding process becomes increasingly expensive as the action dimensionality grows. It would be valuable to include an experiment evaluating how Q-FAT performs in such high-dimensional settings.\n•\tI couldn’t find details on how the demonstration datasets were collected in the appendix. Please indicate whether they were obtained from human demonstrations.", "limitations": "Will Q-FAT still perform well when applied to experiments with inadequate demonstration data? Providing additional experiments or commentary on suboptimal performance would strengthen the paper.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n•\tThe elimination of action quantization in transformer-based policies is an important problem in robotic control systems.\n•\tThe author clearly explains the motivation for adopting GIVT and the Mode-tracking algorithm. The paper is well-structured.\n•\tThe experimental results in simulation are strong.\n\nWeaknesses: \n•\tAlthough the experimental results are state-of-the-art, the paper lacks evaluation on real-world applications. This is a notable gap, especially given that practical utility in real-world robotic control is more important than performance in simulation and that other baselines include such experiments.\n•\tThe authors mention that mode tracking “introduces a minor overhead compared to the transformer forward pass,” but an inference time comparison experiment is missing. Including such an analysis would help substantiate this claim. \n•\t The novelty is somewhat limited, as the GIVT model architecture is adopted from another domain, and the mean-shift algorithm is not new.", "was_revised": true, "final_justification": "After reviewing the author's recent responses, I'm willing to increase my score.", "rebuttal": "Thanks for your thorough evaluation and insightful feedback on our work. We agree that demonstrating practical utility and providing further analysis is important. To address the concerns raised, within the short period of the rebuttal, we have conducted new experiments and will add several key clarifications to the manuscript.\n\n-  **Autonomous Driving Experiments**: We have conducted new experiments on the nuScenes autonomous driving benchmark [Caesar et. al, 2020]. Our preliminary results show that Q-FAT achieves performance competitive with the state-of-the-art VQ-BeT baseline, demonstrating that our method is effective in complex, real-world-adjacent scenarios. Here is the table with preliminary results from our initial experiments (lower is better):\n\n    | Method       | L2 Error (m) | Collision Rate (%) |\n    | :----------- | :----------- | :----------------- |\n    | Diffusion    | 0.96         | 0.44               |\n    | VQ-BeT       | **0.73**         | **0.29**               |\n    | Q-FAT (ours) | 0.75         | 0.37               |\n\n    We plan to add this experiment to our revised manuscript.\n\n-  **Inference Time Analysis**: We have performed an inference time comparison between the sampling algorithms on a trained model from the Kitchen environment, where the model has 4 mixture components. Here is the table of the results:\n\n    | Method        | Time (milliseconds) |\n    | :------------ | :------------------ |\n    | GMM Sampling  | 1.061 ± 0.125   |\n    | Mode Sampling | 1.550 ± 0.314   |\n\n    The performance was tested on a 16 GB Macbook Pro CPU. We will add a more comprehensive study in the revised manuscript across the different environments.\n\n- **Demonstration Data Collection**: We apologize for the omission. The datasets used are standard benchmarks in the field, collected via a mix of human teleoperation and scripted expert policies, consistent with the methods used for our baselines [Chi et. al 2023, Lee et. al 2024]. We will add a dedicated paragraph in the appendix clarifying the source and collection method for each environment's demonstration data. All the data are expert level trajectories, since we are in a behavioral cloning setup.\n\n---\n### References\n[Chi et. al 2023] Diffusion Policy: Visuomotor Policy Learning via Action Diffusion  \n[Lee et. al 2024] Behavior Generation with Latent Actions  \n[Tschannen et. al 2024] JetFormer: An Autoregressive Generative Model of Raw Images and Text"}, {"review_id": "x3HYOlULUt", "reviewer": "Reviewer_Bc6V", "summary": "The paper proposed that Quantization-Free Autoregressive Action Transformer (Q-FAT) replaces the discrete-token step of prior behavior models (e.g., BeT, VQ-BeT) with a decoder-only transformer that directly regresses continuous actions by outputting the means, variances, and logits of a Gaussian-mixture model at every time step. The paper claimed this could preserve the geometry of the action space and remove a separate VQ-VAE stage. To produce smooth yet multimodal trajectories, the authors introduce two generic post-processing techniques: temperature-based variance scaling and a mode-tracking mean shift. Experiments on nine tasks from five simulated‑robotics suites show higher success rates and comparable or better behavioural‐entropy diversity than diffusion‑policy and discrete‑token transformer baselines.", "questions": "1. What is the diversity and quality of the trajectories used for training? (I assume they are all expert-level.)\n\n2. It would be beneficial to include some discussions comparing VQ-based approaches in the context of offline learning for decision-making (e.g., https://openreview.net/forum?id=cA77NrVEuqn, https://openreview.net/forum?id=pQsllTesiE, https://openreview.net/forum?id=n9lew97SAn).", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. The authors conduct experiments across nine tasks, including proprioceptive and image-based environments, both conditional and unconditional, each with 1000 episodes.\n2. Ablation studies are performed regarding the effects of sampling and mixture count.\n3. By omitting the VQ-VAE quantization step used in BeT/VQ-BeT, the proposed method preserves the geometry of continuous controls and simplifies the behavioral cloning pipeline. Empirically, it demonstrates better performance compared to VQ-BeT, which includes the VQ step. Additionally, the GMM parameterization serves as a principled continuous extension of discrete tokens and provides analytic likelihoods, an advantage absent in diffusion policies.\n\nWeaknesses:\n1. The concept of removing the VQ step to maintain the granularity of the action space is not novel (e.g., https://openreview.net/forum?id=btpgDo4u4j). The core transformer approach originates from GIVT, and sampling techniques such as variance scaling are already established in diffusion models. Nevertheless, the presented method does demonstrate performance improvements over prior work, such as VQ-BeT, highlighting the effectiveness of increased granularity in learned actions.\n\n2. While the VQ step indeed introduces additional training costs, it offers benefits in terms of learning robust representations, particularly when the dimensionality of the action space is high. Increased action-space dimensionality may arise in two scenarios: (1) when the action space itself is inherently large, or (2) when learning meaningful action chunks or temporal abstractions becomes critical. To substantiate the claim in the paper (that VQ-based methods limit expressiveness or accuracy in fine-grained control scenarios), a more detailed comparison between methods with VQ (e.g., VQ-BeT) and without VQ (the proposed method) in the described experimental settings would be beneficial.", "was_revised": true, "final_justification": "The proposed method achieves strong experimental performance by effectively maintaining granularity in continuous action spaces. Additionally, as highlighted by the authors during discussions, the method demonstrates clear potential for extension to high-dimensional action spaces. Evaluating the performance of Q-FAT on more complex distributions in high-dimensional settings would be a valuable direction for future work.", "rebuttal": "Thanks for your thorough evaluation and insightful feedback on our work. We will expand the discussion in our paper to include VQ-based approaches in the context of offline learning for decision-making. We address your questions and concerns below:\n\n- **The diversity and quality of the trajectories used for training**: In this work, we are in a behavioral cloning setup (Section 2) where there is no reward signal. This is consistent with the data used with the baseline methods used [Chi et. al 2023, Lee et. al 2024]. Thus, it is standard to assume that the demonstrations are expert level.\n\n- **Vector quantization offers benefits in terms of learning robust representations**: We agree that VQ-based methods could be effective, and our work does not dismiss that. Our central argument is that the VQ step breaks the continuous structure of the action space, which can limit expressiveness, and introduces significant complexity and optimization challenges. These include the need for advanced tricks to handle non-differentiable operations and instabilities during training. Furthermore, it complicates the pipeline by requiring a separate, two-stage training process first for the VQ-VAE and then for the policy. Q-FAT eliminates these issues entirely. We explained an argument why Q-FAT can be more effective (Section 4.1, lines 186-191) than VQ-based methods for fine-grained control. As for GMMs, they universally approximate any distribution given enough mixture components [Palataniotis et al. 2017]. Furthermore, recent work [Tschannen et. al 2024] has demonstrated scaling GMMs to high dimensional image modeling with an autoregressive transformer (in a 128-dimensional latent space with 1k mixture components), demonstrating high sample fidelity. This indicates that scaling GMMs to complex distributions in high-dimensional spaces, such as action spaces, is straightforward.\n\n---\n### References\n[Chi et. al 2023] Diffusion Policy: Visuomotor Policy Learning via Action Diffusion  \n[Lee et. al 2024] Behavior Generation with Latent Actions  \n[Palataniotis et al. 2017] Gaussian Mixtures and Their Applications to Signal Processing  \n[Tschannen et. al 2024] JetFormer: An Autoregressive Generative Model of Raw Images and Text"}]

Author rebuttal:

Meta-review: The paper introduces Q-FAT, a quantization-free autoregressive transformer that directly parameterizes a GMM over actions, simplifying the imitation-learning pipeline while achieving strong performance across diverse simulated robotics tasks. Reviewers highlight clear motivation, thorough experiments (nine tasks), and practical sampling strategies (variance scaling and mean-shift) that reduce trajectory jitter, with two reviews recommending accept and a third moving to borderline accept after rebuttal. The rebuttal also provides preliminary nuScenes results and a concrete inference-time breakdown supporting efficiency claims. The review ratings are all positive. Thus, the AC decides to accept the submission.
Decision: Accept (spotlight)