Pretraining a Shared Q-Network for Data-Efficient Offline Reinforcement Learning

TL;DR

The paper addresses data efficiency in offline RL by introducing a plug-and-play pretraining scheme that initializes a shared Q-network with a backbone $h_\varphi$ and two heads, $g_\psi$ for transition prediction and $f_\theta$ for Q-value estimation. It trains the backbone via a next-state regression objective $L_{pre}$ before fine-tuning in standard offline RL, enabling seamless integration with existing algorithms. The authors analyze the approach using the projected Bellman equation, showing that pretrained representations expand the column space $C(H_\varphi)$, increase the rank of the latent space, and tighten the bound on the Q-value estimation error, which is corroborated by lower proxy Q-errors. Empirically, the method enhances performance and data efficiency across D4RL, Robomimic, V-D4RL, and ExoRL, achieving strong results even with limited data (e.g., 10% of the dataset) and demonstrating robustness across varying data qualities and distributions.

Abstract

Offline reinforcement learning (RL) aims to learn a policy from a static dataset without further interactions with the environment. Collecting sufficiently large datasets for offline RL is exhausting since this data collection requires colossus interactions with environments and becomes tricky when the interaction with the environment is restricted. Hence, how an agent learns the best policy with a minimal static dataset is a crucial issue in offline RL, similar to the sample efficiency problem in online RL. In this paper, we propose a simple yet effective plug-and-play pretraining method to initialize a feature of a Q-network to enhance data efficiency in offline RL. Specifically, we introduce a shared Q-network structure that outputs predictions of the next state and Q-value. We pretrain the shared Q-network through a supervised regression task that predicts a next state and trains the shared Q-network using diverse offline RL methods. Through extensive experiments, we empirically demonstrate that our method enhances the performance of existing popular offline RL methods on the D4RL, Robomimic and V-D4RL benchmarks. Furthermore, we show that our method significantly boosts data-efficient offline RL across various data qualities and data distributions trough D4RL and ExoRL benchmarks. Notably, our method adapted with only 10% of the dataset outperforms standard algorithms even with full datasets.

Figures (15)

• Figure 1: Overview of the pretraining framework. Our approach decomposes the original $Q$-network into two core architectures: a shared backbone network that extracts the representation $z$ from the concatenated state-action input $(s,a)$, and two shallow head networks for learning the transition model and estimating $Q$-values, respectively.
• Figure 2: Reduced approximation error through the expanded column space of $H_\varphi$. In linear approximation, the true value function $Q^\pi$ may lie outside the column space of $H_\varphi$. The projected Bellman equation addresses this by projecting $Q^\pi$ onto its closest representation $\Pi Q^\pi$ within the column space of $H_\varphi$.
• Figure 3: The Rank of the latent space in the $Q$-network during training. We compare the rank of the latent representations between vanilla TD3+BC and TD3+BC+Ours using 512 samples. Our method consistently maintains a higher latent-space rank, indicating a richer feature representation and reduced approximation error.
• Figure 4: Learning curves of TD3+BC. We represent the normalized scores of the vanilla TD3+BC (blue) and TD3+BC (orange) with our pretraining method, respectively. The vertical red dashed lines indicate the transition between the pretraining and main training phases. After pretraining, TD3+BC with our method rapidly surpasses the vanilla baseline by a significant margin.
• Figure 5: Average success rates on the Robomimic benchmark. We compare the baseline methods without pretraining (blue) against those augmented with our pretraining approach (orange) over three seeds. In seven out of eight tasks, our method yields a substantial improvement in success rate across both environments.