Learning Reward–Cost Balance in Safe RL via Score-Based World Models

Safe reinforcement learning (Safe RL) seeks to optimize long-term performance while ensuring adherence to safety constraints. However, most existing approaches address safety in a simplified manner, typically by linearly combining rewards and costs, which provides limited guidance when safety and performance interact in complex, nonlinear ways. We present USB-RL (Unsupervised Score-Balanced Reinforcement Learning), a model-based framework that learns implicit safety–performance preferences directly from experience. Our approach infers a monotone partial-order score through unsupervised pairwise comparisons of long-horizon outcomes, capturing nuanced trade-offs without relying on manually tuned cost weights. The learned score guides model-based policy optimization by dynamically balancing safety and performance, enabling flexible and adaptive multi-step planning in imagination-based control. Across diverse safety benchmarks, USB-RL achieves strong returns while substantially reducing safety violations, demonstrating stable and interpretable safety–performance trade-offs.

AI agents often learn by trial and error, but in real-world tasks such as robot navigation or autonomous driving, mistakes can be unsafe. Safe reinforcement learning aims to help agents achieve their goals while avoiding harmful actions, but many existing methods handle this trade-off by simply subtracting a fixed safety penalty from the reward. This can be too rigid when success and safety interact in complex ways, for example when a robot must move efficiently while avoiding obstacles. We propose USB-RL, a method that teaches an agent to learn its own safety--performance score from experience, without requiring humans to label which behaviors are better. The agent compares long-term outcomes, learns that higher reward should be preferred and higher safety cost should be avoided, and uses this learned score inside a predictive ``world model'' that imagines possible futures before acting. During planning, unsafe imagined actions are filtered out first, and the remaining safe options are ranked by the learned score. In robot-control and driving benchmarks, USB-RL achieves strong task performance while reducing safety violations compared with prior methods. This work may help build learning agents that make more balanced decisions in settings where both achievement and safety matter.