Weighted Implicit Q-Learning Critic

• The paper introduces a weighted implicit Q-learning critic that leverages expectile-weighted regression to safely improve offline RL policies exclusively from logged data.
• It employs a dual-phase update where the value function is optimized via asymmetric regression and the Q-function is updated with a Bellman backup to avoid evaluating out-of-support actions.
• Empirical results on offline RL benchmarks demonstrate state-of-the-art performance, validating the method's stability, sample efficiency, and alignment with constrained optimization theory.

A weighted implicit Q-learning critic is a core component of recent advances in offline reinforcement learning (RL), enabling safe policy improvement exclusively from logged data. The technical innovation centers on fitting a Q-function and an associated value function using expectile-weighted regression, which both restricts learning to the observed action distribution and induces a tunable optimism mechanism for policy improvement. The critic never evaluates out-of-support actions, thus mitigating severe distributional shift. Extensions generalize this approach to arbitrary weighting functions, connect it to behavior-regularized actor-critic paradigms, and derive its structure from constrained optimization principles.

1. Mathematical Foundations of Weighted Implicit Critics

The original implicit Q-learning (IQL) formalism treats the value function $V_\phi(s)$ as a statistic of the random variable $X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$, where $\mu$ is the empirical behavior distribution at state $s$ (Kostrikov et al., 2021). Instead of the mean ($\tau = 0.5$), IQL employs the asymmetric least-squares expectile:

$$m_\tau = \operatorname*{argmin}_m~\mathbb{E}_{x}\left[ |\tau - \mathbb{I}\{x < m\}| (x - m)^2 \right]$$

Here, $\tau \in (0,1)$ interpolates between mean and max; $\tau\to1$ recovers the sample maximum. In practice, $\tau \in [0.7, 0.95]$ is used to select higher-value actions in the data.

The critic loss for $V_\phi$ on data $X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$0 is

$X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$1

This asymmetric regression increases $X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$2 towards higher-Q actions in the dataset, making the value function optimistic relative to the mean.

After fitting $X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$3, the Q-function is updated via a standard Bellman regression:

$X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$4

No maximization or out-of-support action evaluation is performed at any stage of the critic update (Kostrikov et al., 2021).

2. Generalization to Weighted Critic Objectives

The weighted critic update can be generalized as

$X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$5

where $X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$6 is derived from a convex function $X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$7 through $X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$8, with $X_s := \{ Q(s,a)\ |\ a \sim \mu(\cdot|s) \}$9 being the minimizer of the value fitting loss $\mu$0 (Hansen-Estruch et al., 2023). This forms the basis for behavior-regularized actor-critic algorithms with a closed-form implicit policy.

Different weighting schemes correspond to different forms of behavior regularization:

$\mu$1	Weight $\mu$2	Induced policy form
Expectile	$\mu$3	Supports soft improvement, high-$\mu$4 bias
Quantile	$\mu$5	Focus on top quantile actions
Exponential	$\mu$6	Soft-actor (AWR-like), temperature controlled

As the weighting function accentuates high-$\mu$7 actions, the implicit actor diverges further from $\mu$8, increasing exploitation at the cost of behavior regularization (Hansen-Estruch et al., 2023).

3. Policy Implication and Connection to Implicit Actors

At the minimum of $\mu$9, the corresponding implicit policy $s$0 is

$s$1

This form emerges irrespective of the specific convex $s$2. The weight $s$3 thus governs the trade-off between adhering to the behavior policy and exploiting high-value actions. This result (Thm. 4.1 in (Hansen-Estruch et al., 2023), KKT construction in (He et al., 2024)) guarantees Bellman-consistency under the induced policy:

$s$4

The identity explains why weighted regression schemes are intrinsic to implicit Q-learning and its generalizations (He et al., 2024).

4. Constrained Optimization Perspective and AlignIQL

Recent theory formalizes implicit Q-learning’s policy induction as an infinite-dimensional constrained optimization, termed the "Implicit Policy-Finding" (IPF) problem (He et al., 2024). The IPF seeks:

$s$5

Here, $s$6. The resulting optimal policy has the weighted-behavior form:

$s$7

where $s$8 are dual variables learned to enforce normalization and alignment constraints. Specializing to $s$9 yields exponential weights, directly linking to the Advantage-Weighted Regression recipe. AlignIQL-hard implements these duals with parameterized networks, while soft variants use a scalar relaxation parameter $\tau = 0.5$0:

$\tau = 0.5$1

All variants adjust both actor extraction and (optionally) critic regression with these weights, thus strictly aligning the induced policy and critic (He et al., 2024).

5. Algorithmic Implementation and Stability Considerations

Weighted implicit Q-learning critics are implemented via alternating updates of $\tau = 0.5$2 (expectile regression) and $\tau = 0.5$3 (weighted Bellman backup), with soft target network updates for stability. Canonical pseudocode for IQL critic update (Kostrikov et al., 2021):

1. Sample minibatch $\tau = 0.5$4 from $\tau = 0.5$5
2. Compute TD errors $\tau = 0.5$6
3. Update value: $\tau = 0.5$7
4. Update Q: $\tau = 0.5$8
5. Target Q update: $\tau = 0.5$9

For general weighted objectives, $$m_\tau = \operatorname*{argmin}_m~\mathbb{E}_{x}\left[ |\tau - \mathbb{I}\{x < m\}| (x - m)^2 \right]$$0 is weighted as above (Hansen-Estruch et al., 2023), and analogously in AlignIQL with policy-derived $$m_\tau = \operatorname*{argmin}_m~\mathbb{E}_{x}\left[ |\tau - \mathbb{I}\{x < m\}| (x - m)^2 \right]$$1 (He et al., 2024). Stabilization strategies include double Q-networks, gradient clipping, and using minibatch sampling without TD-error prioritization.

6. Significance in Offline RL and Empirical Performance

The weighted implicit critic addresses the central challenge of offline RL: balancing policy improvement with avoidance of extrapolation error from out-of-distribution action evaluation (Kostrikov et al., 2021). By only using dataset-supported actions and controlling optimism via $$m_\tau = \operatorname*{argmin}_m~\mathbb{E}_{x}\left[ |\tau - \mathbb{I}\{x < m\}| (x - m)^2 \right]$$2 or the weighting function, the critic achieves both conservatism and sample-efficient improvement.

Empirical studies demonstrate that this architecture achieves state-of-the-art performance on standard offline RL suites such as D4RL, including robust generalization, improved stability, and strong results on challenging sparse-reward domains (Kostrikov et al., 2021, He et al., 2024). AlignIQL further increases empirical alignment between the actor and critic and shows enhanced hyperparameter robustness and superior results on Antmaze and Adroit tasks compared to the original IQL and IDQL (He et al., 2024).

7. Extensions, Relations, and Theoretical Insights

Weighted implicit Q-learning critics have yielded extensive theoretical and algorithmic extensions:

• Generalized weighting through arbitrary convex losses unifies expectile, quantile, and exponential (AWR-like) weighting (Hansen-Estruch et al., 2023).
• The constrained optimization (IPF) perspective establishes that the weighted regression form is not heuristic, but the unique KKT solution for a natural regularized policy extraction problem, and that the induced weights must appear in both actor fit and (optionally) in Q backup for optimal alignment (He et al., 2024).
• The diffusion-behavior models in IDQL and AlignIQL enable expressive, multimodal implicit actors by sampling actions from learned diffusion models and weighting them with the critic-derived $$m_\tau = \operatorname*{argmin}_m~\mathbb{E}_{x}\left[ |\tau - \mathbb{I}\{x < m\}| (x - m)^2 \right]$$3 (Hansen-Estruch et al., 2023, He et al., 2024).

This line of research resolves conceptual ambiguities in earlier IQL variants, solidifies the necessity and optimality of weighted regression for both critic and actor, and provides both practical and theoretical guidance for scalable, stable offline RL.