RaPO: Retention-Aware Policy Optimization Fixes Catastrophic Forgetting in MLLMs

Introduction

In the LLM/MLLM domain, verifier-rewarded reinforcement fine-tuning (RFT) such as GRPO has shown significant advantages in eliciting complex reasoning. Due to its on-policy nature, RFT is believed to have stronger anti-forgetting properties than supervised fine-tuning (SFT). However, whether standard RFT can effectively resist catastrophic forgetting in challenging visual continual learning (VCL) settings — such as rehearsal-free class-incremental and domain-incremental learning — remains unexplored.

Problem: Trajectory-Level Drift Agnosticism in GRPO

The HKU team (Prof. Yizhou Yu's group) conducted a systematic study using Qwen2-VL-2B on ImageNet-R with a 10-task, 5-shot class-incremental setup. Results (Figure 1a) show GRPO outperforms SFT overall, but its accuracy still drops sharply as tasks accumulate, indicating unresolved stability-plasticity imbalance.

To diagnose the cause, the authors analyzed the KL divergence between candidate rollouts (generated by the current Actor policy) and the frozen Anchor policy from the previous task. Even when rollouts achieve identical high task rewards, their KL divergence from the Anchor policy varies widely (Figure 1b):

Low-drift trajectories : preserve the previous task's reasoning patterns and feature distributions while learning new classes.

High-drift trajectories : match the new task label but undergo drastic parameter-space and representation shifts, abandoning historical knowledge.

Standard GRPO computes group-relative advantage solely from task reward, assigning similar credit to both low- and high-drift trajectories. This Trajectory-level Drift Agnosticism lets the optimizer drift toward high-drift regions, causing catastrophic forgetting. Causality was confirmed by two GRPO variants: keeping only low-drift rewards reduced forgetting; keeping only high-drift rewards worsened it. Hard thresholding is unstable, so the authors propose turning drift into a continuous reward signal.

Method: RaPO (Retention-aware Policy Optimization)

RaPO introduces two plug-and-play components without adding extra parameters or architectural changes.

1. Retention Reward: Bounded Reward Reshaping via Trajectory-Level Drift

Instead of a global KL regularization term (which cannot differentiate rollouts within a group), RaPO maps each trajectory's drift to a dense scalar Retention Reward that participates in the group-relative advantage calculation before the advantage is computed.

Drift is averaged over trajectory length for fair comparison; a one-sided max(·,0) truncation prevents reward hacking.

Smaller drift yields reward near 1; larger drift decays toward 0. The total reward combines task reward and retention reward:

When task rewards are similar, lower-drift trajectories receive higher advantage and are more strongly reinforced. λ balances the two terms; the mechanism works with any task-specific reward (classification accuracy, detection IoU, etc.).

2. Cross-Task Advantage Normalization (CTAN)

At task boundaries, GRPO's per-group reward standard deviation σ collapses (near zero when most rollouts are correct), inflating advantages; at the start of a new task, reward drops and variance spikes, compressing advantages. This oscillation destabilizes training.

CTAN replaces the instantaneous batch standard deviation with a cross-task persistent EMA standard deviation σ̂:

β is close to 1 (e.g., 0.99), so σ̂ evolves slowly, smoothing out boundary spikes. The advantage denominator uses σ̂, which is saved and carried across tasks, providing a stable optimization scale (Figure 2).

Experiments: Five Visual Continual Learning Benchmarks

1. Class-Incremental Image Classification (CIL)

On ImageNet-R, ImageNet-A, TinyImageNet, and CUB-200 (Qwen2-VL-2B, 10–20 tasks, 5-shot), SFT with classic continual learning regularizers (L2, EWC, LwF) fell below 40% accuracy. Standard GRPO achieved 74.67% last accuracy on ImageNet-R (10 tasks) but with 20.02% forgetting. RaPO raised final accuracy to 85.92% and reduced forgetting to 4.69% . On harder ImageNet-A, RaPO improved accuracy by 7.24% and cut forgetting by 8.72% over GRPO.

2. Class-Incremental Object Detection (CIL)

Using Qwen2-VL-7B on COCO 2017 (5- and 10-task splits), RaPO achieved Box AP of 19.31% and 19.12%, improving over GRPO by 4.67% and 4.82%, while forgetting dropped to 1.39% and 1.37%. This shows RaPO's retention mechanism works with structured detection rewards computed via Hungarian matching.

3. Video Classification & Domain-Incremental Learning (DIL)

On UCF-101 and Kinetics-200 (video classification), RaPO maintained best accuracy after long task sequences. In domain-incremental settings (DomainNet for classification, Pascal cross-domain detection), RaPO outperformed all baselines; on Pascal cross-domain detection it reached 37.18% average AP, surpassing the joint-training SFT upper bound of 35.40%.

4. Qualitative Analysis: Preserving Structured Reasoning Trajectories

After full multi-task training, standard GRPO's <think> chains confuse early classes with later ones (e.g., misidentifying a candle as a bald eagle, or missing bicycles in detection). RaPO's reasoning chains retain fine-grained discriminative features for old tasks, demonstrating that the retention reward not only protects final labels but also consolidates the underlying structured visual reasoning trajectories.

Conclusion

RaPO refutes the assumption that RFT is inherently immune to catastrophic forgetting. It pinpoints trajectory-level drift agnosticism as the core bottleneck and reconstructs RFT's credit assignment and variance normalization with simple mathematics. RaPO integrates seamlessly into existing PPO/GRPO frameworks and provides a solid algorithmic foundation for building lifelong-learning adaptive embodied and visual systems. Code is open-sourced at https://github.com/LMMMEng/RaPO; paper available at https://arxiv.org/abs/2605.09640 (NeurIPS 2026).