Imagination-Limited Q-Learning for Offline Reinforcement Learning

1 Introduction

Offline Reinforcement Learning (RL) [Lange et al., 2012; Fujimoto et al., 2019] is designed to learn optimal policies purely from a static dataset previously collected by an unknown policy (behavior policy). By eliminating the need for online interaction with environments, it offers dual benefits. On the one hand, it can mitigate the expensive costs [Gu et al., 2017] and potential risks [Sallab et al., 2017] associated with trial-and-error learning in real-world applications; on the other hand, it can be leveraged to enhance the generalization ability and scalability of RL models when the logged data is massive and diverse. However, the offline learning paradigm unavoidably incurs distributional shifts [Levine et al., 2020] of state-action visitation between the learned policy and the behavior policy, which makes it difficult to correctly assess out-of-distribution (OOD) action-values. Especially, overoptimistic estimates may even invalidate the learned policy.

∗Corresponding author

To address this challenge, two main classes of technical routes are commonly employed in the model-free approaches [Prudencio et al., 2023]. 1) Policy constraint: It usually explicitly restricts the gap between the learned and behavior policy. The batch-constrained Q-learning (BCQ) [Fujimoto et al., 2019] was devised to restrict the action space via adding perturbations on a state-conditioned behavior model. Kumar et al. [Kumar et al., 2019] proposed BEAR to reduce maximum mean discrepancy between the learned policy and the behavior one. Subsequently, different methods corresponding to other metrics have been proposed, such as KL divergence for BRAC [Wu et al., 2019] and mean squared error for TD3-BC [Fujimoto and Gu, 2021; Srinivasan and Knottenbelt, 2024], which similarly seeks to steer the policy closer to actions in the dataset. However, these learned models are limited to the neighborhood of the behavior policy hindering their performance, especially when the dataset is collected by poor policies. 2) Value regularization: It aims to utilize value regularizations to suppress OOD action-values. The conservative Q-learning (CQL) [Kumar et al., 2020] was designed to penalize the expectation of OOD action-values, thus mitigating optimistic estimates outside the dataset. Kostrikov et al. [Kostrikov et al., 2022] introduced implicit Q-learning (IQL) to estimate value function through expectile regression to implicitly depress OOD action-values. The MCQ [Lyu et al., 2022] was proposed to regularize OOD action-values with the maximum behavior value. And the OAC-BVR [Huang et al. ] was developed to regard the difference between the Q-function and the behavior value as a regularization term. While these methods effectively limit OOD optimism, they also introduce uncontrollable bias into Q-value estimates, as illustrated in Fig. 1(a). Especially, the Q-values under CQL exhibit noticeable bias of pessimism over all MuJoCo tasks, as shown in Fig. 1(c).

To mitigate value bias while maintaining appropriate restrictions on over-optimism, we propose an ImaginationLimited Q-learning (ILQ) method. The insight of our method is straightforward: For in-sample state-action pairs, we adopt the standard Bellman backup based on in-sample transitions to estimate their values. For OOD state-action pairs, instead of blindly applying value regularizations to suppress their action-values, we envision what the values would be without any restrictions. The one-step bootstrapped values under an imaged dynamics model would be reasonable estimations, ideally approximating the ground truth when the imaged model closely matches the environment. However, errors in the model fitting are inevitable, and optimistic estimates may still exist. Therefore, we need to further consider how to appropriately limit the imagination values. We employ the maximum behavior value as the ceiling of the imagined one. Specifically, if the imagined value is less than the maximum value, it is maintained; otherwise, the maximum behavior value is applied. Figure 1(b) illustrates the intuition behind our method. This design ensures a more reasonable evaluation with appropriate constraints on OOD actions, thereby avoiding unnecessary value suppression and improving the generalization ability of the learned policy.

We validate the effectiveness of ILQ both theoretically and empirically. We prove that the policy evaluation Bellman operator of ILQ is a contraction operator under tabular Markov Decision Processes (MDPs), ensuring convergence. Particularly, we analyze the action-value gap between the fixed point obtained by our method and that obtained by the Bellman optimality equation, demonstrating that the error bound of OOD action-values can reach the same order of magnitude as in-distribution ones. Empirically, our method maintains reasonably optimistic Q-values compared to conservative Qlearning (Fig. 1(c)) and achieves state-of-the-art performance across a wide range of tasks in the standard benchmark.

2 Background

Reinforcement Learning (RL) is commonly modeled as a Markov Decision Process (MDP), characterized by a tuple (S, A, r, P, ρ0, γ) [Sutton and Barto, 2018], where S is state space, A is action space, r : S ×A → [-rmax, rmax] is reward function, P : S×A×S → [0, 1] is transition dynamics, ρ0 is probability distribution of initial states, and γ ∈ [0, 1) is discount factor. The RL agent takes actions on the environment according to its policy, defined as π : S ×A → [0, 1].

During the learning process, evaluating the expected return of a state-action pair (s, a) under a policy π, called the actionvalue (or Q-value) Q(s, a), is essential. Theoretically, it can be obtained by Eπ[∑∞t=0 γ t r(st, at) | s0 = s, a0 = a]. In practice, it is usually approximated by minimizing the Bellman residual E[(Q(s, a) - (T Q)(s, a))2 ], where T is the Bellman optimality operator defined as

Generally, a delayed approximator Q- is applied in the above target T Q for training stability [Mnih et al., 2015].

2.1 Offline Reinforcement Learning

In offline RL, the agent is no longer allowed to interact with the environment and learns policies exclusively from a limited static dataset D, which is typically gathered from an unknown behavior policy β. The dataset D is usually represented by a set of transition tuples {(s, a, r, s′ )}. The goal of the agent remains to maximize the expectation of cumulative rewards. Without online correction, erroneously optimistic estimates of OOD action-values are inevitable [Levine et al., 2020]. These biases can cause the learned policy to favor incorrect OOD actions, potentially leading to it failure.

The representative value regularization method is CQL [Kumar et al., 2020], which adds penalties for OOD actionvalues to the standard Q-value update objective, as follows:

where µ(·|s) is a distribution to produce OOD actions, αCQL is a hyperparameter to adjust the degree of conservatism.

3 RelatedWork

3.1 Model-free Offline RL

Recent advancements in offline RL have focused on addressing the challenges posed by OOD actions. Importance sampling methods [Nachum et al., 2019] have been developed to correct evaluation under distributional shifts, but they often suffer from high variance. Policy constraint methods [Kumar et al., 2019; Wu et al., 2021; Fujimoto and Gu, 2021; Li et al., 2023; Chen et al., 2024] are employed to limit the deviation of learned policies; however, their performance tends to degrade when the behavior policy is suboptimal.

Our method aligns more closely with value regularization approaches. For instance, CQL [Kumar et al., 2020] and CSVE [Chen et al., 2023] directly penalize OOD actionvalues, effectively controlling over-optimism but often resulting in overly pessimistic estimates. Methods like MCQ [Lyu et al., 2022] and OAC-BVR [Huang et al., 2024] attempt to relax restrictions by assigning maximum behavior values or behavior values, respectively, to OOD action values. However, this introduces uncontrolled value bias for OOD actions. Although MCQ additionally employs policy constraint weighting to mitigate this bias in practical implementations, it offers limited theoretical guarantees for OOD action-value estimates. In contrast, our proposed method preserves more honest estimates of OOD action-values within the limitation of maximum behavior values, and offers theoretical guarantees for its value estimates.

3.2 Model-based Offline RL

Model-based methods aim to enhance collected datasets by generating synthetic trajectories using learned dynamics models. Various strategies have been proposed to effectively leverage these synthetic data, including uncertainty quantification [Ovadia et al., 2019; Kidambi et al., 2020; Yu et al., 2020; Diehl et al., 2021], conservative value estimation [Yu et al., 2021], representation balancing [Lee et al., 2021], and adversarial learning [Bhardwaj et al., 2023]. In contrast, our proposed ILQ avoids trajectories generation, relying solely on the dynamics model to produce one-step subsequent states and rewards of in-sample states for estimating imagined OOD action-values. While ILQ utilizes the dynamics model, it remains fundamentally a model-free learning framework and circumvents challenges of error accumulation associated with longer trajectories in model-based methods.

4 Imagination-Limited Q-learningMethod

We start by elucidating our novel Imagination-Limited Bellman (ILB) operator in Subsection 4.1. Subsequently, we elaborate on its practical implementation details and theoretical analysis in Subsection 4.2 and 4.3, respectively. And the Imagination-Limited Q-learning (ILQ) algorithm is ultimately summarized in Subsection 4.4.

4.1 Imagination-Limited Bellman Operator

In online RL, researchers typically use the Bellman optimality operator Eq. (1) to evaluate policies. However, in offline settings, the absence of online corrections makes policy evaluation highly susceptible to OOD over-optimism [Levine et al., 2020] under the standard operator. Existing value regularization methods primarily focus on directly restricting the OOD action-values [Kumar et al., 2020; Lyu et al., 2022; Chen et al., 2023; Huang et al., 2024], which introduces uncontrollable bias in value estimates and lacks theoretical guarantees for OOD actions.

We argue that establishing reasonable estimates for out-ofdistribution state-actions should take precedence, followed by the imposition of suitable restrictions, rather than directly employing value regularization. To achieve this goal, we introduce a novel Imagination-Limited Bellman operator, defined as follows.

Definition 1. The Imagination-Limited Bellman (ILB) operator is defined as

where β is the behavior policy,

and

are the imagined value and its limitation, respectively. The ̂ P is the empirical transition kernel, ̂ r is the empirical reward function, δ is a hyperparameter with a small absolute value, and Supp(·) means support-constrained on the dataset.

Here is the insight behind the proposed ILB operator. For an in-sample state-action pair (s, a), i.e., β(a | s) > 0, we have its corresponding transition (s, a, r, s′ ) in the dataset, allowing us to apply the standard Bellman operator without any obstacles. However, for an out-of-sample state-action pair (s, aoos ), the standard Bellman backup cannot be applied solely due to the absence of its successor state and reward. To address this, one could utilize empirical dynamics model to predict the next state ̂ s′ and reward ̂ r and obtain the imagination value yQ img as Eq. (3), which provides a relatively accurate approximation for an OOD state-action. Nevertheless, it may still result in optimistic estimates because of fitting errors. To tackle this issue, we use the maximum in-distribution action-value Eq. (4) as the upper limit for the imagined value. This design offers dual benefits: First, it maximally maintains the imagined value, reducing estimation bias on OOD actions; second, the maximum behavior value ensures that there is always an in-distribution action-value greater than or equal to the OOD ones, encouraging the policy to more likely favor in-distribution actions during the actor improvement process.

We analyze the ILB operator's properties and demonstrate that the ILB operator exhibits the same γ-contraction property as the standard Bellman operator, ensuring convergence in policy evaluation. The proof is provided in the Appendix.

Theorem 1 (Convergence). The ILB operator defined in Eq. (2) is a γ-contraction operator in the L∞ norm, and Qfunction iteration rule obeying the ILB operator can converge to a unique fixed point.

4.2 Practical Implementation of Imagination and Limitation Value

Imagination Value

We fit the environment dynamics to derive empirical transfer kernel ̂ P and reward function ̂ r in the simplest manner using

where ̂ Tψ stands for both ̂ P and ̂ r for brevity, and is represented by a multivariate Gaussian distribution with parameters ψ practically. We then obtain the imagined value yQ img .

Limitation Value

In fact, the behavior policy β(· | s) in Eq. (4) is unknown and needs to be empirically modeled. In light of the expressiveness of diffusion models [Ho et al., 2020], we fit the behavior policy using a conditional diffusion model. Specifically, it is constructed via a reverse diffusion chain, formulated as

where superscript k denotes the diffusion timestep, a := a0 is the final sampled action, ak , k = 1, · · · ,K - 1, are latent variables, aK ∼ N(0, I) is Gaussian noise. Typically, pω(ak-1 | ak , s) is modeled as a Gaussian distribution N (ak-1 ; µω(ak , s, k), Σω(ak , s, k)) with the covariance matrix Σω(ak , s, k) = βkI and the mean defined as

where βk is the variance schedule, αk := 1 - βk, ¯αk := ∏ki=1 αi, and ξω(·) is the noise prediction network with parameters ω. The conditional diffusion model is optimized by maximizing the evidence lower bound, which can be simplified [Ho et al., 2020] to minimize the following objective

where U is an uniform distribution over {1, · · · ,K}. Similar diffusion behavior modeling methods are also applied in other works [Wang et al., 2023; Hansen-Estruch et al., 2023].

Accordingly, we adopt the limitation value as shown in the following equation:

where M is the number of sampled actions. Due to possible errors in the fitting process, this may result in a deviation between the estimated value Eq. (9) and the true value Eq. (4). In order to offset this gap, we introduced the hyperparameter δ in definition Eq. (2), typically set to a small absolute value.

4.3 Theoretical Analysis

We now theoretically discuss the action-value gap between the fixed point in Theorem 1 and the Bellman optimality value. Before proceeding further, we make some commonly used assumptions about the reward function [Huang et al., 2024, Assumption 1].

1. The reward function is bounded, i.e., |r(s, a)| ≤ rmax. Actually, this is consistent with what is required by its definition r(s, a) : S × A → [-rmax, rmax].
2. Similar to the Lipschitz condition, i.e., |r(s, ˜a1) - r(s, ˜a2)| ≤ ℓ∥˜a1 - ˜a2∥∞, ∀s ∈ S and ∀˜a1, ˜a2 ∈ A, where ℓ is a constant. This requires that the reward function satisfies Lipschitz continuity with respect to actions.

In addition, the error bound assumption between the empirical models and the real ones are required, which is also utilized in both [Kumar et al., 2020] and [Huang et al., 2024]. Suppose the ̂ r and ̂ P are the empirical reward function and empirical transition dynamics, respectively, the following relationships hold with high probability ≥ 1 - ζ, ζ ∈ (0, 1),

where D is the constant related to the dataset size, ζr and ζP are constants related to ζ.

We begin by analyzing the Bellman optimality value gap between the learned policy and behavior policy.

Theorem 2. Suppose Qβ∗ is the fixed point of the supportconstrained Bellman optimality operator. The following gap can be obtained

Accordingly, we could prove the error bound between the imagination value and Bellman optimality value.

Theorem 3. Suppose Qβ∗ is the fixed point of supportconstrained Bellman optimality operator. The gap between the imagination value yQβ∗ img and Qβ∗ has:

Based on above theorems, we can estimate the action-value gap between the fixed point of the ILB operator and Qβ∗ .

Theorem 4 (Action-value gap). Suppose QILB and Qβ∗ denote the fixed point of the ILB operator and supportconstrained Bellman optimality operator, separately. The action-value gap can be bounded as

where ζr, ζP are defined in Eq. (10) and Eq. (11), ϵr, ϵP are defined in Theorem 2, δ is defined in the ILB operator.

According to Theorem 4, we conclude that error bounds for in-sample and out-of-sample actions are of the same magnitude O(rmax/(1-γ)2 ). This result aligns with the conclusion of CQL [Kumar et al., 2020] within the support region. All proofs are detailed in the Appendix.

4.4 The ILQAlgorithm

In deep RL, the Q-function is commonly approximated by a neural network with parameters θ, while the corresponding target network has parameters θ-. As described in the background, it can be optimized by minimizing the temporal difference (TD) loss E(s,a,r,s′ )[(Qθ(s, a) - T Qθ-(s, a))2 ]. Intuitively, under the ILB operator we developed, the corresponding loss function can be constructed as

Since in-sample and out-of-sample state-action pairs have different TD targets, as defined by our ILB operator, a common way is to split the above loss function into a weighted sum of in-sample and out-of-sample components as follows

where η is a trade-off factor, (s, aoos ) refers to the out-ofsample state-action pair. The u(· | s) is any distribution that can produce out-of-sample actions. In practice, we directly treat policy π as the sampling distribution u.

We also incorporate the double network [Hasselt, 2010] to reduce overestimation, which is widely utilized in both online [Mnih et al., 2015] and offline RL [Fujimoto and Gu, 2021]. Therefore, combining the definition of ILB operator Eq. (2), we have the target, for in-sample transition (s, a, r, s′ ), as:

where πϕ is the learned policy represented by a neural network with parameters ϕ and Qθj is the j-th target network. Similarly, for out-of-sample (s, aoos ), the target is formulated as:

where

Here Eq. (20) is obtained by coupling Eq. (9).

During policy improvement, we adopt the same objective as in vanilla SAC [Haarnoja et al., 2018] to optimize the actor network without any complex design, as follows:

where α is a multiplier for the entropy.

Combining above steps, our method is derived, with its pseudo-code presented in Algorithm 1.

Algorithm 1 Imagination-Limited Q-Learning (ILQ)

• Require: The offline dataset D, number of iterations N, discount factor γ, target network update rate τ , trade-off factor η, and offset parameter δ.
• 1: Initialize critic networks Qθ1 , Qθ2 , actor network πϕ, target networks Qθ1 , Qθ2 with θi ← θi, i = {1, 2}.
• 2: // Pre-train the dynamics model and behavior policy
• 3: Train the dynamics model ̂ Tψ(s′ , r | s, a) via (5).
• 4: Train the diffusion model Diffω(· | s) for modeling the behavior policy by optimizing (8).
• 5: // Policy training
• 6: for step n = 1 to N do
• 7: Sample a mini-batch of transitions B = {(s, a, r, s′ )} from dataset D.
• 8: Compute target value for (s, a) in B as (17).
• 9: Sample OOD actions conditioned on sates in B via πϕ, and calculate target value according to (18) via the pretrained behavior policy and dynamics model.
• 10: Update parameters θi, i = 1, 2 for each critic network via minimizing (16).
• 11: Update actor ϕ via (21).
• 12: Update the target networks

• 13: end for

5 Experiments

In this section, we empirically validate the effectiveness of our method ILQ. 1) We demonstrate the superiority of ILQ over existing methods by comparing performance across a series of tasks. 2) We conduct sensitivity analyses on the hyperparameters involved in ILQ, confirming the stability of the proposed method. 3) We then perform ablation experiments on both imagination and limitation components to verify their impacts. 4) We also delve into the Q-value estimations to further validate the effectiveness of the two components designed; details are in the Appendix due to space constraints.

5.1 Experimental Settings

We evaluate ILQ on the D4RL [Fu et al., 2020] benchmark. The commonly used domain is GymMuJoCo "-v2", including halfcheetah, hopper, and walker2d tasks at four levels: random (r), medium (m), medium-replay (mr), and mediumexpert (me). We also assess ILQ on Maze2D "-v1" domain, which offers three layouts with two reward types, i.e., umaze (u), umaze-dense (ud), medium (m), medium-dense (md), large (l), and large-dense (ld). In addition, comparisons on several adroit "-v0" tasks are conducted. Due to page limitations, descriptions of compared algorithms and extra experimental results have been relocated to the Appendix. Details of hyperparameters settings and implementation specifics are also provided in the Appendix to ensure reproducibility.

5.2 Performance Comparison

In comparing ILQ with state-of-the-art methods on MuJoCo tasks, as shown in Table 1, ILQ performs significantly better than policy constraint methods (BCQ [Fujimoto et al.,

2019], UWAC [Wu et al., 2021], One-step [Brandfonbrener et al., 2021], TD3+BC [Fujimoto and Gu, 2021], and OAP [Yang et al., 2023]) in a wide range of random- and mediumlevel tasks. This is because policy constraint methods restrict the learned policy to be within a neighborhood of the behavior policy. Although TD3-BST [Srinivasan and Knottenbelt, 2024] introduces fine-grained weighting on the constraints of TD3+BC to enhance performance, it remains inferior to ILQ overall. The value regularization methods (CQL [Kumar et al., 2020], MCQ [Lyu et al., 2022], CSVE [Chen et al., 2023], OAC-BVR [Huang et al., 2024]) show higher scores in average. ILQ continues to show performance beyond them in 11/12 tasks. DTQL [Chen et al., 2024] integrates implicit value regularization and policy constraints to enhance performance, but it still falls short compared to ILQ, especially in halfcheetah tasks.

According to Table 2, neither policy constraint approaches nor value regularization approaches performed well enough on Maze2D tasks. Although ROMI-BCQ [Wang et al., 2021] achieves advanced performance on maze2d-u by utilizing a reverse dynamics model, it performs mediocrely on other tasks. Diffuser [Janner et al., 2022] and PlanCP [Sun et al., 2023] leverage diffusion models to improve their planning capabilities in maze2d-u, while still lag behind ILQ on maze2dm and maze2d-l, further demonstrating stitching abilities of ILQ. The experimental results on Adroit tasks demonstrate that our method continues to outperform others, highlighting its applicability across different tasks.

5.3 Parameter Study

Offset parameter δ

To assess impacts of parameter δ, we conduct sensitivity analyses by varying δ across {-2,-1,-0.5, 0, 0.5, 1, 2}. We evaluate the performance on halfcheetah-me, hopper-mr, hopper-me, and walker2d-m, where each experiment was run over 3 random seeds. Figure 2 shows that the score curve remains stable as δ changes. Overall, performance is slightly lower when δ is negative compared to when it is positive. Notably, a negative δ always ensures a safe learned policy. However, when δ is positive, meaning the limit exceeds the maximum behavior value a little bit, there may exists policy failure, as illustrated in Fig. 2(c). These indeed validate our argument that overly restricting OOD values could inhibit potential performance gains, and taking the maximum behavior value (δ = 0) as a ceiling of the imagined value is consistently a solid choice.

Trade-off Factor η

To evaluate the impact of the trade-off factor η, we conduct sensitivity analyses by varying η around its optimal value.

According to results in Fig. 3(a) and 3(c), ILQ achieves robust performance for all variations of η around 0.9, making it a typically safe choice for medium and medium-replay tasks. This suggests that ILQ should place greater trust in in-sample value estimates when evaluating policies. In medium tasks, the Q-value of OOD actions generated by the learned policy rarely reach the maximum behavior value. Consequently, assigning too high a weight to these OOD action-values, i.e., a smaller η, may lead to misplaced trust in OOD actions and ultimately cause policy failure, as illustrated in Fig. 3(c). In medium-expert tasks, the in-sample data comprises a mixture of medium and expert levels, and the value of the OOD actions generated by the learned policy can approach the maximum behavior value. Therefore, the value estimate needs to be more balanced between in-sample and out-of-sample. In this case, the optimal value of η is usually around 0.6 for all medium-expert tasks. The overall results in Fig. 3 show that ILQ maintains stable performance when η varies around its optimal parameter.

5.4 Ablation Study

To understand the contribution of each component in our OOD target action-value Eq. (18), we conduct an ablation study. This study evaluates the impact of removing either the imagination component or the limitation value from the target value. All experiments are run over 3 random seeds. More results can be found in the Appendix.

Without Imagination

In this part, we assess the performance of ILQ without the imagination component yQ img , indicating solely the maximum behavior value is considered as the regularization target. The results are illustrated in Fig. 4. As shown in Fig. 4(a) and (b), the performance degrades significantly and the curve of normalized score exhibits a very oscillatory behavior. This is because directly using the maximum behavior value as the target values for OOD actions introduces uncontrollable bias and ultimately impairs policy improvement. This demonstrates the necessity of the imagination component for providing calibrated target values under the limitation.

Without Limitation

Here we exclude the limitation component yQ lmt from the target value, relying solely on yQ img for learning. This approach is intended to evaluate the importance of the limitation component. In one specific case, Fig. 5(a), performance increased when relying solely on the imagination value, indicating that the imagination component can sometimes provide highly reliable guidance. However, in most tasks, performance dropped significantly, with some policies in hopper-m task Fig. 5(b) collapsing completely. This underscores the importance of the limitation component in preventing the incorrectly optimistic estimates. These studies suggest both of the two components play critical role on OOD estimates.

6 Conclusion

In conclusion, the Imagination-Limited Q-learning (ILQ) method effectively mitigates bias of value estimations by maintaining reasonable evaluations of OOD action-values within appropriate limits. Specifically, it utilizes a dynamics model to help generate imagined values and capping these with the maximum behavior values for OOD actions, while standard target values for in-distribution ones. Theoretical analysis confirms the convergence of ILQ and demonstrates that the error bound between estimated and optimal values for OOD actions is comparable to that for in-distribution actions, thereby enhancing performance improvements. Empirical results show that ILQ achieves state-of-the-art performances on a wide range of tasks in the D4RL benchmark. We hope this work can provide new insights into the value estimates for offline RL.

Acknowledgments

This work was partially supported by the National Science Foundation of China (62072185 and U1711262). The Appendix is available at https://github.com/LiuWH-AI/ILQ.