Extending QVIRL to Continuous Action Spaces — Six Candidate Directions

Self-authored study document mapping out six candidate directions for extending ValueWalk from discrete to continuous action spaces. Parent: thesis-hub. This is the reference to return to when picking the thesis’s main technical direction.

Why this matters

QVIRL’s speed comes from two closed-form tricks (the max-mean linearization for the inverse Bellman equation, and the Lu et al. mean-field softmax approximation) — both defined over a discrete action set $\mathcal{A}$. Going continuous breaks discreteness in exactly three places: the partition function in the likelihood, the argmax in the inverse Bellman equation, and policy extraction at the end. Every direction below is defined by which of these it fixes and how, and — critically — whether Lemma 1 (closed-form propagation of Q-uncertainty into the reward posterior) survives. That survival check is the viability test for any candidate.

The six directions (summary)

1. Structured Q-functions with analytic action integrals (NAF-style) — force $Q$ to be quadratic in the action (à la NAF), so both the argmax and the partition integral become exact closed forms (Gaussian integral). Cleanest fit with QVIRL’s existing machinery; Lemma 1 not only survives but simplifies. Cost: unimodal in the action (mixture-of-quadratics extension recovers multimodality at moderate extra cost).
2. Amortized maximization — an actor network — import the DDPG/TD3 pattern: train an actor ${\pi }_{\varphi }(s)$ to stand in for the argmax, and reuse it as an importance-sampling proposal for the partition function. Lemma 1 survives verbatim. Main risk: actor suboptimality and the max-mean approximation both bias rewards upward in the same direction, compounding rather than cancelling. This is the “expected”/baseline direction — build it first regardless of what becomes the main contribution.
3. The partition integral as a point-process object (LGCP connection) — recognize that $Z(s)=\int e^{\beta Q^{\ast }(s,a)}da$ is the intensity integral of a log-Gaussian Cox process, and import closed-form variational bounds from spatial statistics (Lloyd et al.). Produces a genuine lower bound (Jensen, right direction) with analytic pieces under RBF kernels. Doesn’t touch the Bellman-max problem — composable with Direction 1 or 2. The most theory-heavy, most “QVIRL-flavored” option (bespoke closed-form derivation over sampling).
4. Autoregressive action discretization (what VLAs actually do) — factorize the action into per-dimension bins chosen sequentially (Q-Transformer-style), so both closed-form tricks apply per-dimension unchanged. $K^d$ joint problem becomes $d$ sequential $K$-way problems. Catch: a factorized softmax is a different expert model than a joint softmax, which matters for IRL (unlike forward RL) because the likelihood is the expert model — characterizing that bias is the theory contribution needed here.
5. Score matching / diffusion — dissolving the partition function entirely — the identity ${\nabla }_a\log {\pi }_E(a|s)=\beta {\nabla }_a{Q}^{\ast }(s,a)$ means the partition function’s gradient vanishes, so score matching (à la Q-Score Matching) never needs to compute it. Removes Bottleneck 1 completely, but a score-matching objective isn’t a log-likelihood, so the ELBO becomes a generalized/Gibbs posterior rather than orthodox variational Bayes — must be invoked openly, not smuggled in. Independently useful even if not the main objective: diffusion samplers help with contrastive-action sampling (Direction 2) and multimodal risk-averse policy extraction.
6. Soft-Bellman reframing — replace the hard max with the soft value $V_{soft}(s)=\frac{1}{\beta }\log \int e^{\beta Q(s,a)}da$ (the IQ-Learn / SAC move), unifying Bottlenecks 1 and 2 into one logsumexp, estimated via a SAC-style inner variational policy. Listed last because it requires the reparameterization trick QVIRL’s authors specifically engineered around for stability — viability hinges on an open empirical question (does that instability reproduce here?). Forfeits the closed-form Lu/max-mean tricks; more a different method occupying the same problem than a strict extension.

Comparison table

Direction	Fixes partition (B1)	Fixes argmax (B2)	Lemma 1 / closed-form KL	ELBO still a true bound?	Main risk	Effort
1. Quadratic Q (NAF)	Exactly	Exactly	Survives, simplifies	Yes	Unimodality in actions	Low–medium
2. Actor + proposals	Approx. (importance sampling)	Approx. (amortized)	Survives verbatim	Sampling noise in likelihood	Compounding underestimation biases	Low
3. LGCP bound	Yes (Jensen/kernel integrals)	No — pair with 1 or 2	Untouched	Yes — bound direction correct	Bound looseness; ignores num.–denom. correlation	Medium–high (theory)
4. Autoregressive bins	Reduces to discrete Lu, per dim	Reduces to discrete max, per dim	Survives per dimension	Yes, for the factorized model	Factorized ≠ joint expert model	Medium
5. Score matching	Dissolved (no normalizer)	No — pair with 1 or 2	Untouched	No — generalized posterior	Losing orthodox Bayesian semantics	High (theory)
6. Soft-Bellman	Unified into $V_{soft}$	Unified into $V_{soft}$	Forfeited	Nested bound, but sampled	Reparameterization instability	High

Suggested plan

• Phase 1 — baseline: Direction 2 (actor + importance-weighted contrastives). Exercises the full pipeline at MuJoCo scale; everything else gets evaluated against it. Deliverable: working continuous QVIRL + a corrected, importance-weighted continuous version of the Lu approximation (the current Appendix B.2 sampling stopgap silently estimates the wrong quantity — no importance weights).
• Phase 2 — main contribution, pick one:
  • Direction 1 for a safe, fully-closed-form method with a clean expert-model story (mixture-of-quadratics as stretch goal for multimodality).
  • Direction 3 for the theoretical centerpiece — a derived bound in the group’s GP idiom (note: Osborne, a ValueWalk author, also co-authored the key Lloyd et al. LGCP reference), composable with Phase 1’s actor.
  • Direction 4 if the robotics/VLA deployment story is the emphasis.
• Phase 3 — payoff experiment: calibration and risk-aversion in continuous action spaces — does the posterior stay calibrated post-extension (checked against MCMC ground truth in low dimensions, mirroring QVIRL’s own validation), and does the $\mu -k\sigma$ risk-averse apprentice measurably avoid uncertain regions on a continuous-control safety task? Direction 5’s energy-guided sampler is a natural drop-in for multimodal risk-averse policy extraction even if score matching isn’t the main training objective.

Two standing cautions before committing to a direction: (1) run the Lemma 1 viability test for any candidate mechanism before investing further — several attractive tricks fail there, quietly; (2) the reparameterization-instability question is load-bearing for Directions 5 and 6 and cheap to test early — a one-week experiment reproducing (or not) the instability in a continuous setting would close or open a third of the design space.

Key reading by direction

• D1 (NAF): Gu, Lillicrap, Sutskever & Levine, Continuous Deep Q-Learning with Model-based Acceleration (ICML 2016); Amos, Xu & Kolter, Input Convex Neural Networks (ICML 2017) if quadratic proves too rigid.
• D2 (actor): Lillicrap et al., DDPG (2015); Fujimoto et al., TD3 (2018); Garg et al., IQ-Learn (NeurIPS 2021); Van de Wiele et al., Q-Learning in Enormous Action Spaces via Amortized Approximate Maximization (2020).
• D3 (LGCP): Lloyd, Gunter, Osborne & Roberts, Variational Inference for Gaussian Process Modulated Poisson Processes (ICML 2015); Møller, Syversveen & Waagepetersen (1998); John & Hensman, Large-Scale Cox Process Inference (ICML 2018); Hennig, Osborne & Kersting, Probabilistic Numerics (2022).
• D4 (autoregressive bins): Chebotar et al., Q-Transformer (CoRL 2023); Metz et al., SDQN (2017); Tang & Agrawal (AAAI 2020); Tavakoli, Pardo & Kormushev, Action Branching Architectures (AAAI 2018); RT-2 (Brohan et al. 2023), OpenVLA (Kim et al. 2024), FAST (Pertsch et al. 2025).
• D5 (score matching/diffusion): Psenka, Escontrela, Abbeel & Ma, Q-Score Matching (ICML 2024); Florence et al., Implicit Behavioral Cloning (CoRL 2021) — read early; Chi et al., Diffusion Policy (RSS 2023); π0 (Black et al. 2024); Wang et al., Diffusion-QL (2022); Hansen-Estruch et al., IDQL (2023); Lu et al., Contrastive Energy Prediction (ICML 2023); Bissiri, Holmes & Walker (JRSS-B 2016); Knoblauch, Jewson & Damoulas, Generalized Variational Inference (2019).
• D6 (soft-Bellman): Ziebart’s PhD thesis (2010); Haarnoja et al., Soft Q-Learning (2017) and SAC (2018); Garg et al., IQ-Learn (2021).

My status / decision

TBD — which direction(s) to pursue, and why. Discuss Phase 1 vs. jumping straight to a Phase 2 candidate with Bajgar.