AdamO: A Collapse-Suppressed Optimizer for Offline RL

Abstract: Offline reinforcement learning (RL) can fail spectacularly when bootstrapped temporal-difference (TD) updates amplify their own errors, driving the critic toward extreme and unusable Q-values. A key counterintuitive insight of this work is that collapse is not only a property of the backup rule or network architecture: optimizer dynamics themselves can directly trigger or suppress instability. From a control-theoretic viewpoint, we model offline TD learning as a feedback system and analyze Adam-based critic updates. This yields a necessary and sufficient condition for stability of the induced local update dynamics: within the regime we analyze, these dynamics are stable if and only if the spectral radius of the corresponding update operator is strictly below one. Further analysis suggests that standard Adam updates can inadvertently distort the parameter geometry, motivating explicit orthogonality constraints to prevent TD error amplification. To this end, we propose AdamO, an Adam-based optimizer with a decoupled orthogonality correction regulated by a strict task-alignment budget. We prove that this design theoretically guarantees worst-case task safety and preserves Adam's continuous-time dissipative dynamics. Empirically, AdamO is broadly compatible with diverse offline RL baselines, improving stability and returns across a broad suite of benchmarks.

• The paper introduces AdamO, an Adam-based optimizer that prevents value collapse in offline RL by enforcing local stability conditions through an orthogonality correction.
• It presents a control-theoretic framework that analyzes critic dynamics via spectral decomposition, establishing necessary and sufficient conditions for exponential convergence.
• Empirical evaluations on D4RL benchmarks demonstrate that AdamO outperforms standard optimizers, achieving up to 123% improvement in challenging tasks.

AdamO: A Collapse-Suppressed Optimizer for Offline RL

Introduction and Motivation

Offline reinforcement learning (RL) is fundamentally challenged by instability arising from the "deadly triad" of bootstrapping, function approximation, and off-policy training, which leads to unbounded value divergence and collapse of the critic. Traditional remedies—policy constraints, value regularization, target networks—address symptoms rather than the root causes within the optimization dynamics. This paper rigorously demonstrates that optimizer dynamics themselves can precipitate or prevent collapse in offline RL, presenting an optimization-theoretic framework that leads to a principled diagnosis and remedy.

Control-Theoretic Analysis of Critic Dynamics

The paper adopts a control-theoretic lens, modeling offline TD learning as a feedback system where the TD error recursively interacts with its own bootstrapped predictions. By linearizing the update operator and analyzing Adam-based critic updates, the paper derives a closed-form, second-order recurrence for the TD error. A key theoretical finding is the identification of a necessary and sufficient local stability condition: the spectral radius $\rho(\mathsf{A}(\eta)) < 1$ of an augmented update operator $\mathsf{A}(\eta)$ is both necessary and sufficient for exponential convergence of the critic TD error. This criterion is made precise in Theorem 4.1, which formalizes the precise relationship between optimizer dynamics and collapse in offline RL.

Figure 1: Convergence versus collapse of the critic loss in TD3+BC, illustrating the connection to contraction versus expansion of the spectral radius of the TD update operator $\mathsf{A}(\eta)$.

Notably, this perspective transcends prior collapse analyses limited to SGD, highlighting the qualitative differences introduced by Adam’s adaptive, higher-order dynamics.

Spectral Decomposition of Collapse: Stability Knobs

The stability criterion is further unpacked via spectral analysis of the TD update operator. The Hurwitz condition for operator $\mathsf{S}$—determined by the sign of the real parts of its eigenvalues—delineates regimes of stable contraction and unstable self-excitation. The spectral structure is shown to decompose into two independently controllable effects:

• Bootstrapped Scale Term: Controlled by spectral norm constraints and input normalization/clipping, which bound amplification due to the bootstrapping of targets.
• Feature-Conditioning Term: Quantifies the deviation of parameter-induced geometry from isometry and is governed by the (near-)orthogonality of the weight matrices.

This decomposition motivates a stability recipe: scale can be bounded by standard normalization and spectral constraints, while the geometry requires an explicit mechanism for orthogonality in the parameter space.

AdamO: Optimizer-Level Orthogonality Correction

Building on the above analysis, the paper introduces AdamO, an Adam-based optimizer featuring a decoupled, task-aligned orthogonality correction on selected weight blocks. The correction is designed to:

• Prevent amplification of TD errors by directly controlling parameter-side geometry (orthogonality defect).
• Avoid contaminating Adam’s adaptive moment estimates, which would otherwise distort optimization if orthogonality is enforced via a loss regularizer.
• Limit interference with first-order task descent through a per-layer conflict budget $\tau$ and scale control $\kappa$.

The update rule applies a normalized orthogonalization drift after the Adam step, scale-matched, and budget-constrained to never undermine primary convergence unless a limited amount of misalignment is permitted.

Theoretical Guarantees

AdamO’s design is matched by rigorous theoretical guarantees:

• Conflict-Free Mode ($\tau=0$): Provably ensures non-inferiority in next-step task loss compared to Adam, provided standard smoothness and stepsize conditions.
• Budgeted Mode ($\tau>0$): Theoretical upper bounds are derived quantifying the worst-case single-step degradation as a function of conflict budget and orthogonality strength.
• Hamiltonian Perspective: In continuous time, AdamO preserves Adam's dissipative Lyapunov/Hamiltonian structure, with the orthogonality module introducing at most bounded additional drift.

Hyperparameter admissible ranges for $\kappa$ are explicitly characterized.

Empirical Evaluation

Extensive empirical analysis is conducted on D4RL MuJoCo Locomotion, AntMaze, Adroit, and Kitchen offline RL benchmarks, spanning diverse difficulty regimes and reward structures. AdamO is tested as a drop-in replacement for Adam in six representative baselines (TD3+BC, IQL, ReBRAC, ACTIVE, PARS, SQOG) and compared with other stability-oriented optimizers and normalization schemes. Experimental findings are summarized as follows:

• Consistent Performance Gains: AdamO outperforms Adam and other first-order optimizers across domains, most decisively in difficult, collapse-prone tasks (AntMaze, Adroit-human/cloned).
• Collapse Suppression: In challenging regimes where standard methods exhibit value explosion or divergence, AdamO stabilizes the learning curves and preserves high returns.

Figure 2: Comparison of computational overhead and wall-clock runtime across different optimizers, showing that AdamO incurs only minor additional cost.

Figure 3: Radar charts show that AdamO expands the normalized score frontier across domains and algorithmic backbones.

Strong quantitative improvements include: over 100% average return increase versus Adam in AntMaze, and 123% relative improvement on Locomotion tasks in the 10k data regime.

Ablations and Mechanistic Analyses

Sensitivity studies on correction scale $\kappa$ and conflict budget $\mathsf{A}(\eta)$0 demonstrate that AdamO’s stability and efficacy are robust to reasonable choices of these parameters, with critical thresholds matching the theoretical predictions.

Figure 4: Sensitivity analysis of TD3+BC across different $\mathsf{A}(\eta)$1 values, demonstrating the critical threshold effect on stability.

Figure 5: Ablation on conflict budget $\mathsf{A}(\eta)$2, confirming that moderate, controlled budgets balance stability and learning performance.

Figure 6: Critic loss curves for IQL on 20 D4RL tasks, showing AdamO eliminates loss spikes and deviation.

Spectral analyses confirm that AdamO maintains the Hurwitz condition on the local operator in practice, aligning the leading eigenvalues’ real part with stability.

Figure 7: Training dynamics—the leading eigenmode of the TD operator and its real/imaginary components—demonstrate AdamO constrains self-excitation and prevents divergence.

Implications and Future Directions

AdamO is the first optimizer-level intervention in offline RL that (a) explicitly targets the spectral mechanism of value collapse, (b) operates at the optimizer rather than as an auxiliary regularization or architectural change, and (c) is justified by necessary and sufficient stability conditions. The methodology opens several avenues:

• Theoretically principled optimizer interventions in other bootstrapping-heavy learning systems.
• Optimization-geometric regularization as a substitute for fragile heuristic regularizers.
• Fine control of training stability under non-stationarity and distribution shift.
• Extension to high-dimensional, data-scarce, and multi-task settings.

One limitation remains an inherent trade-off: infinitesimal step sizes, while maximizing stability, are computationally prohibitive. This suggests that a full solution to stability must blend optimizer and architectural innovations as well as adaptive approximation of the relevant spectral quantities in practice.

Conclusion

This paper presents a principled optimizer-centric solution to value collapse in offline RL. By formalizing the spectral roots of critic instability and introducing AdamO, an optimizer with provable collapse suppression, this work positions optimizer dynamics as a first-class control point for RL algorithm design. Empirical results confirm that optimizer-level spectral control yields substantive performance gains, and the theoretical analysis lays a foundation for future methods that go beyond traditional regularization and architecture-driven stability fixes.

References

[AdamO: A Collapse-Suppressed Optimizer for Offline RL, (2605.01968)]