Do You Really Need to Pretrain Q-Functions for Online RL Fine-Tuning?

Overview§

This paper investigates whether pretraining the Q-function is beneficial when fine-tuning a pretrained policy in online reinforcement learning. The authors find that naive Q-function pretraining provides little to no improvement over random initialization due to a fundamental mismatch: the pretrained Q-function targets the pretrained policy, but online fine-tuning converges to a different Q-function. They propose Initialization via Policy Ensemble (IPE), a method that trains multiple diverse policies on offline data and uses their pooled rollouts to initialize the Q-function, leading to significant improvements in fine-tuning performance.

Methodology§

Problem Setup§

Consider an offline dataset $\mathcal{D} = \{(s_i, a_i, r_i, s'_i)\}$ collected by some behavior policy. We have a pretrained base policy $\pi_{\text{base}}$ (e.g., from behavioral cloning). The goal is to fine-tune $\pi_{\text{base}}$ via online RL using a Q-function $Q_\theta$. The standard approach initializes $Q_\theta$ by pretraining on $\mathcal{D}$ using the Bellman equation: $$ \mathcal{L}_{\text{pretrain}} = \mathbb{E}_{(s,a,r,s')\sim\mathcal{D}}\left[\left(Q_\theta(s,a) - (r + \gamma \mathbb{E}_{a'\sim\pi_{\text{base}}(\cdot|s')}[Q_{\bar{\theta}}(s',a')]\right)^2\right] $$ During online fine-tuning, the Q-function is updated with Bellman backups using the current policy $\pi$.

The Mismatch§

Define $Q^{\pi_{\text{base}}}$ as the true Q-function of $\pi_{\text{base}}$ and $Q^{}$ as the optimal Q-function (the target of online RL). The pretrained Q-function approximates $Q^{\pi_{\text{base}}}$, but online fine-tuning converges to $Q^{}$ (or the Q-function of the final policy). These can be very different, so initialization from $Q^{\pi_{\text{base}}}$ does not accelerate convergence to $Q^{*}$. The authors prove that even offline value maximization (e.g., CQL) does not bridge this gap.

IPE: Initialization via Policy Ensemble§

IPE trains $K$ diverse policies $\{\pi_i\}_{i=1}^K$ on the offline data using different random seeds, behavioral cloning, or offline RL. Then, it collects a pooled set of rollouts (simulated or from the dataset) from all policies. The Q-function is initialized by regressing onto the ensemble of Q-values: $$ \mathcal{L}_{\text{IPE}} = \mathbb{E}_{(s,a)\sim\mathcal{D}_{\text{pool}}} \left[ \left( Q_\theta(s,a) - \frac{1}{K} \sum_{i=1}^K \hat{Q}^{\pi_i}(s,a) \right)^2 \right], $$ where $\hat{Q}^{\pi_i}$ are estimated via Monte Carlo returns or Bellman evaluation. This provides a more informative initialization that covers diverse behaviors, reducing the mismatch.

Implementation Details§

Training Procedure§

1. Offline phase: Train $K$ policies (e.g., using BC or TD3+BC) with different seeds. For each policy, estimate its Q-function via TD learning or Monte Carlo returns. 2. Pooled rollouts: Collect transitions from all policies (either by simulating or using saved transitions). 3. Q-function initialization: Train $Q_\theta$ on the pooled data using IPE loss. 4. Online fine-tuning: Run standard online RL (e.g., SAC) with $Q_\theta$ initialization, freezing the base policy or allowing it to be fine-tuned.

Code Snippet (PyTorch-style)§

# Offline: train K policies and get their Q estimates
policy_ensemble = [train_offline_policy(seed=i) for i in range(K)]
Q_targets = torch.zeros(len(pooled_dataset), 1)
for i, (s, a) in enumerate(pooled_loader):
    Q_i = [policy.Q(s, a) for policy in policy_ensemble]
    Q_targets[i] = torch.mean(torch.stack(Q_i))

# IPE initialization
Q_theta = QNetwork()
optimizer = Adam(Q_theta.parameters(), lr=3e-4)
for epoch in range(N_initialization_epochs):
    for batch in pooled_loader:
        s, a, target = batch
        pred = Q_theta(s, a)
        loss = F.mse_loss(pred, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# Online fine-tuning (e.g., SAC)
for step in range(fine_tune_steps):
    # ... standard SAC updates using Q_theta as initial Q

Results§

IPE outperforms naive Q-function pretraining on continuous control benchmarks (HalfCheetah, Hopper, Walker2d) with an average 1.26x improvement in fine-tuning returns. It also matches or exceeds the performance of more complex offline-to-online algorithms.

Key Takeaways§

• Q-function pretraining is often not beneficial due to policy mismatch.
• IPE provides a simple, effective initialization by ensembling policies.
• The method requires no architectural changes and minimal computational overhead (K policies can be trained in parallel).