Lifted bellman linear programming improves offline reinforcement learning training

Lifted Bellman Linear Programming for Offline Reinforcement Learning

Machine LearningArtificial Intelligence

Summary

Offline reinforcement learning teaches AI systems to make decisions using only pre-collected data, which is tricky because it usually requires complex steps to estimate future rewards accurately. The authors introduce a new approach that enforces Bellman optimality directly within the given data, avoiding unstable regression targets and the need for target networks. Their method lifts the problem into a joint space of value functions, ensuring the learned values respect dataset constraints better, especially when environment dynamics are predictable. They implement a practical version with neural networks that performs well without complex stabilization techniques and uses fewer computing resources.

What this means in practice

  • For machine learning engineers: Build more stable offline RL critics that do not require target networks or EMA updates, reducing complexity and computation.
  • For robotics developers: Develop robot control policies from pre-recorded data without needing complex off-policy corrections, especially in deterministic environments.

Authors

Hyukjun Yang, Jongchan Park, Narim Jeong, Donghwan Lee

Abstract

Offline reinforcement learning (RL) typically trains a critic by minimizing a regression loss against bootstrapped value targets stabilized by target networks with exponential moving average (EMA) updates. Multi-step targets incorporate behavior-policy actions and therefore require off-policy correction. We instead impose in-sample Bellman optimality on the critic through inequality constraints. We formulate the Lifted Bellman Linear Program (LBLP), which lifts the linear programming characterization of Bellman optimality to the joint $(Q,V)$ space so that every constraint involves only state-action pairs in the dataset. Its unique minimizer is the in-sample optimal pair, and constraints along $K$-step segments of dataset trajectories leave this minimizer unchanged for any rollout policy and horizon. Under deterministic dynamics, this minimizer lies between the best dataset return and the optimal value. Relaxing the constraints into hinge penalties recovers the same solution above a finite penalty coefficient in the tabular case. Approximate Lifted Bellman Unconstrained Minimization (ALBUM) implements this relaxation with neural networks and detaches the $K$-step rollout targets by stop gradient. Its objective contains no squared regression onto bootstrapped targets, so it can be trained without target networks or EMA updates. Under deterministic dynamics, the LBLP solution is a stationary point of the detached update under a coefficient condition independent of $γ$ and $K$, and the inequality constraints allow discounted returns along dataset trajectories to serve as lower bounds without off-policy correction or action chunking. On OGBench, ALBUM uses a single critic with a Gaussian policy, matches the average performance of FQL, and is comparable to recent action-chunking methods, while using the fewest parameters and the least peak GPU memory among all compared methods.