Reflect-then-Correct: Rebalancing Task Optimization for Generalizable Meta-Reinforcement Learning via Distributional Value Error Reduction

Meta-Reinforcement Learning (Meta-RL) faces significant challenges in non-parametric settings, where vastly different return scales across diverse tasks cause severe gradient interference. Existing categorical solutions attempt to normalize these scales but often fail due to rigid discretization and quantization errors. To address this, we propose Reflect-then-Correct (RTC), a framework that models meta-values using Sinkhorn divergence. By treating distributions as adaptive floating particles, RTC achieves a geometry-aware alignment of distinct meta-task structures. However, while Sinkhorn updates harmonize gradients, they introduce statistical bias via sampling estimation. RTC overcomes this issue by "reflecting'' on the temporal accumulation of Bellman inconsistencies through a recursive error model and "correcting'' the optimization via adaptive importance weights, which prioritize more accurate transitions for meta-value estimation. We provide theoretical guarantees for this reweighting strategy and demonstrate that RTC outperforms existing baselines on the challenging Meta-World ML-10 and ML-45 benchmarks.

Modern AI systems often struggle to adapt quickly when faced with highly diverse tasks, particularly when rewards and objectives vary substantially. This challenge is especially prominent in Meta-Reinforcement Learning (Meta-RL), where an agent must learn how to adapt efficiently across many environments. Existing methods often simplify task differences by forcing information into fixed categories, which can reduce learning accuracy and lead to unstable training. We propose a new approach, called Reflect-then-Correct (RTC), that enables AI systems to represent task information more flexibly by modeling values as adaptive distributions rather than rigid categories. RTC identifies inconsistencies that accumulate during learning and prioritizes more accurate transitions for meta-value estimation. We provide theoretical support for this reweighting strategy and demonstrate experimentally that RTC improves both learning efficiency and generalization on challenging robotic tasks compared with existing methods.