[Paper] A Single-Loop Bilevel Deep Learning Method for Optimal Control of Obstacle Problems

Source: arXiv - 2601.04120v1

Overview

This paper tackles the notoriously hard problem of optimal control for obstacle problems—situations where a physical or simulated system must stay above (or below) a moving “obstacle” while minimizing a cost. Traditional solvers rely on fine mesh discretizations and nested optimization loops that quickly become computationally prohibitive in high‑dimensional or irregular domains. The authors introduce a single‑loop bilevel deep‑learning framework that replaces the mesh‑based sub‑solvers with neural networks, dramatically cutting runtime while preserving solution quality.

Key Contributions

• Mesh‑free bilevel formulation: Uses neural networks to represent both the state (solution of the PDE) and the control, eliminating the need for costly mesh generation.
• Constraint‑embedding networks: Network architectures are designed to automatically satisfy the obstacle constraints, so the optimizer never steps into infeasible regions.
• Single‑Loop Stochastic First‑Order Bilevel Algorithm (S2‑FOBA): A novel training algorithm that removes the inner‑outer loop structure typical of bilevel problems, enabling end‑to‑end gradient‑based learning.
• Convergence analysis: Provides theoretical guarantees for S2‑FOBA under mild smoothness and bounded variance assumptions, without requiring a unique lower‑level solution.
• Extensive empirical validation: Demonstrates comparable or better accuracy than classical finite‑element methods on benchmark distributed‑control and obstacle‑control tasks, with up to an order of magnitude speed‑up.

Methodology

1. Problem encoding

   • The obstacle optimal control problem is expressed as a bilevel optimization: the lower level solves a variational inequality (the PDE with obstacle constraints), while the upper level minimizes a performance functional over the control variables.

2. Neural surrogate models

   • Two neural networks, $ \mathcal{N}\theta $ (state) and $ \mathcal{M}\phi $ (control), are introduced.
   • The state network is built with a constraint‑embedding layer (e.g., a ReLU or projection onto the feasible set) that guarantees the obstacle inequality is always satisfied.

3. Single‑loop training (S2‑FOBA)

   • Instead of solving the lower‑level problem to optimality at each outer iteration, the algorithm treats the lower‑level optimality condition as a stochastic first‑order residual and updates both $ \theta $ and $ \phi $ simultaneously using unbiased gradient estimators.
   • Mini‑batch sampling of collocation points in the domain provides stochastic estimates of the PDE residual and the objective gradient.
   • A carefully chosen stepsize schedule ensures that the coupled updates converge to a stationary point of the original bilevel problem.

4. Implementation details

   • Mesh‑free collocation points are drawn from simple distributions (uniform or Sobol sequences), making the method trivially scalable to high dimensions.
   • Automatic differentiation frameworks (PyTorch, JAX) compute all required gradients, so the pipeline integrates cleanly with existing deep‑learning toolchains.

Results & Findings

• Accuracy: Across all benchmarks, the neural surrogate achieved ≤ 4 % relative error compared with high‑resolution finite‑element solutions.
• Speed: Because S2‑FOBA avoids repeated solves of the lower‑level PDE, total wall‑clock time dropped to 8–12 % of the classical approach.
• Scalability: Experiments on a 3‑D domain with $10^{6}$ collocation points showed linear memory growth and stable convergence, confirming the method’s suitability for large‑scale problems.
• Robustness: The algorithm converged even when the lower‑level problem had multiple feasible solutions, thanks to the relaxed uniqueness requirement in the theory.

Practical Implications

• Rapid prototyping: Engineers can now embed obstacle‑type constraints (e.g., safety margins in robotics, contact constraints in simulation) directly into neural controllers without hand‑crafting mesh pipelines.
• Edge deployment: Since the trained networks are lightweight inference models, optimal‑control policies can run on embedded devices (micro‑controllers, GPUs) with real‑time latency.
• Design optimization loops: In industries like aerospace or additive manufacturing, where obstacle constraints evolve during design iterations, the single‑loop approach enables continuous re‑optimization without costly re‑meshing.
• Integration with existing ML stacks: The method plugs into PyTorch/JAX, allowing developers to combine it with reinforcement learning, differentiable physics, or meta‑learning pipelines.

Limitations & Future Work

• Assumption of smooth PDE coefficients: The convergence proof relies on Lipschitz continuity; highly discontinuous material properties may degrade performance.
• Sample efficiency: While mesh‑free, the stochastic estimator still needs a relatively large number of collocation points for high‑accuracy PDE residuals, which can be memory‑intensive.
• Extension to time‑dependent obstacles: The current formulation handles static obstacles; handling moving or dynamic obstacles would require recurrent or physics‑informed temporal networks.
• Theoretical gap for nonconvex upper‑level objectives: The analysis guarantees convergence to stationary points, but global optimality remains open for highly nonconvex cost functions.

The authors suggest exploring adaptive sampling strategies, hybrid physics‑informed neural networks, and multi‑level extensions to tackle time‑dependent and stochastic obstacle problems.

Authors

• Yongcun Song
• Shangzhi Zeng
• Jin Zhang
• Lvgang Zhang

Paper Information

• arXiv ID: 2601.04120v1
• Categories: math.OC, cs.LG
• Published: January 7, 2026
• PDF: Download PDF