Prioritized Model Experience Replay
Abstract
Model-based reinforcement learning (MBRL) improves sample efficiency by leveraging learned dynamics models, but often suffers from unstable training due to dynamics model learning mismatch: models are trained on data from historical policies while being queried under the continually updated current policy. This mismatch can cause policy-relevant local model error to remain large even as global prediction error decreases, leading to oscillatory updates. We present a finite-horizon performance analysis that decomposes the policy performance gap into global model error, policy-induced distribution shift, and historical policy mixture effects, showing that minimizing global error alone is insufficient for stable optimization. Motivated by this analysis, we propose Prioritized Model Experience Replay (PMER), a lightweight replay mechanism that prioritizes high-error transitions during dynamics model training. PMER implicitly emphasizes policy-relevant regions without explicit policy distance estimation and integrates seamlessly into Dyna-style MBRL frameworks. Experiments on MuJoCo benchmarks demonstrate improved stability, faster convergence, and higher sample efficiency.
Lay Summary
Reinforcement learning agents improve through trial and error, but collecting experience can be expensive, slow, or risky. Model-based reinforcement learning reduces this cost by using the agent’s real interaction experience to train a virtual environment, so the agent can "practice" in imagination before acting in the real world. However, a mismatch arises because the agent keeps changing: as it improves, it explores new regions of the real environment and collects new experiences, while the virtual environment is often trained uniformly on historical experiences. This can make the virtual environment slow to learn these newly explored regions, so the agent may practice in places where the simulation is locally inaccurate, leading to slow and unstable learning. During virtual environment training, we can identify experiences from these poorly modeled regions using the virtual environment’s own prediction errors and replay them more often. Based on this idea, we propose PMER, which helps the virtual environment become more accurate in the places that matter most, therefore allowing the agent to learn more effectively through virtual trial and error and improve faster in the real environment.