First-order Sobolev Reinforcement Learning
Fabian Schramm, Nicolas Perrin-Gilbert, Justin Carpentier
TL;DR
The paper addresses the gap in off-policy RL where standard TD learning ignores local geometry of the value function. It introduces First-Order Sobolev RL, which enforces first-order Bellman consistency by differentiating Bellman backups through differentiable dynamics and adding gradient-target matching to the critic via a Sobolev loss. The approach yields a Sobolev critic that minimizes $(Q_\phi - y)^2$ plus gradient terms on $\nabla_s Q_\phi$ and $\nabla_a Q_\phi$, and is compatible with Q-learning and actor–critic methods like DDPG and SAC. Empirical results on a 1D toy problem and a differentiable Ant task show faster early learning and more stable value estimates, suggesting practical benefits for gradient-based policy updates and potential for adaptive weighting or line-search techniques to further balance value and gradient consistency.
Abstract
We propose a refinement of temporal-difference learning that enforces first-order Bellman consistency: the learned value function is trained to match not only the Bellman targets in value but also their derivatives with respect to states and actions. By differentiating the Bellman backup through differentiable dynamics, we obtain analytically consistent gradient targets. Incorporating these into the critic objective using a Sobolev-type loss encourages the critic to align with both the value and local geometry of the target function. This first-order TD matching principle can be seamlessly integrated into existing algorithms, such as Q-learning or actor-critic methods (e.g., DDPG, SAC), potentially leading to faster critic convergence and more stable policy gradients without altering their overall structure.