PaperlayerView exact revision on arXiv

Source-linked AI summary

AegisFlow: Training-Free Non-Myopic Path-Safe Guided Flow Matching

Kunpeng Liu, Boshi Zhang, Keyou You

arXiv:2610.39825v1cs.AI

TL;DR

Training-free path-safe guidance methods rely on myopic min-norm controllers that may produce suboptimal trajectories near safety boundaries. AegisFlow incorporates a time-varying barrier into a global optimal-control cost and approximates the HJB equation; it reports a better trade-off between target alignment and constraint satisfaction across image manipulation and robotic planning.

  • Problem

    Recent path-safe guidance methods use myopic min-norm controllers that can yield suboptimal trajectories near safety boundaries, limiting target optimization.

    Source: §1 · p. 2, §3 · p. 5

  • Method

    AegisFlow formulates guided generation as global safety-critical optimal control, integrates a time-varying barrier into the cost functional, and derives a tractable HJB approximation.

    Source: §1 · pp. 2–3, §4.1 · p. 6, §4.2 · p. 6

  • Results

    AegisFlow reports a better trade-off between target alignment and constraint satisfaction across image manipulation and continuous control, including improved identity retention, success rates, trajectory smoothness, and safety.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The evaluations support AegisFlow's effectiveness across high-dimensional image manipulation and safety-critical continuous control tasks.

    Source: §1 · p. 3

    Path-safety of the optimal control under regularity conditions does not automatically guarantee safety of the HJB-approximated control or Algorithm 1 output; a safe guard may be used when necessary.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

Pretrained generative models, such as diffusion and flow matching models, have demonstrated exceptional capabilities in synthesizing high-quality data. However, steering their generative trajectories to satisfy strict constraints--such as identity preservation in image editing or obstacle avoidance in robotics--remains fundamentally challenging. While recent training-free guidance methods enforce path-safety, they rely heavily on min-norm controllers that only account for immediate constraint satisfaction. These myopic approaches often yield suboptimal trajectories that hover dangerously close to the boundary, severely restricting the model's ability to optimize toward its target. In this paper, we propose AegisFlow, a training-free, non-myopic path-safe guidance framework that recasts guided generation as a safety-critical optimal control problem with global scope. By integrating a time-varying barrier function directly into the global cost functional, AegisFlow naturally enforces constraints along the entire generative path. To compute the global optimal control, we derive a tractable, mathematically grounded approximation of the Hamilton-Jacobi-Bellman (HJB) equation. Comprehensive experiments demonstrate that AegisFlow outperforms state-of-the-art methods across text-guided image manipulation and robotic planning, establishing a significantly better trade-off between target alignment and constraint satisfaction. Code is available at https://github.com/LiuKunpeng03/AegisFlow.

1 Introduction

The introduction motivates path-safe guidance as a way to steer pretrained generative trajectories toward targets without violating constraints. AegisFlow replaces myopic corrections with global optimal control, and is reported to improve the trade-off between target alignment and constraint satisfaction across image manipulation and continuous control.

1 Introduction

  • Prior path-safe methods use myopic min-norm control, which can yield suboptimal trajectories near safety boundaries and limit target optimization.Soft guidance lacks hard guarantees, while projection can degrade generation quality.Source: §1 · p. 2
  • AegisFlow frames guided generation as global safety-critical optimal control, integrating a time-varying barrier into the cost functional to account for long-term effects.The authors derive a tractable approximation to the associated Hamilton-Jacobi-Bellman equation; finite-cost solutions satisfy the path constraints.Source: §1 · pp. 2–3
  • The evaluation spans text-guided image manipulation on CelebA-HQ and physical planning, including robotic manipulation.The introduction reports stronger textual alignment and identity retention in image editing, and higher success, smoother trajectories, and fewer safety violations in continuous control than current safe diffusion methods.Source: §1 · pp. 2–3
  • AegisFlow establishes a better trade-off between target alignment and constraint satisfaction across high-dimensional image manipulation and safety-critical continuous control.In image editing it maintains a larger identity-preservation safety margin while improving textual alignment and identity retention; in continuous control it improves success rates and trajectory smoothness while reducing safety violations.Source: §1 · p. 3

2 Related Works

Related work covers training-free image manipulation and control-theoretic approaches to flow matching. The paper distinguishes its globally scoped barrier-cost approach from prior methods that guide or constrain trajectories locally.

2 Related Works

  • Training-free image manipulation commonly inverts an image to noise before guiding the generative trajectory, while other approaches optimize the original image or avoid inversion.The paper focuses mainly on inversion-based methods and leaves application to other approaches for future work.Source: §2 · p. 3
  • Prior flow-matching control methods use PMP or HJB to optimize costs, but regularizers do not provide hard guarantees on trajectory behavior.The paper also reviews per-step projection and control barrier functions with min-norm controllers for enforcing path constraints.Source: §2 · p. 3
  • AegisFlow estimates optimal control from an HJB equation whose cost functional includes a barrier, targeting global optimization while keeping trajectories within a safe set.The related-work discussion contrasts this with per-step methods and with approaches it characterizes as essentially myopic or heuristic.Source: §2 · p. 3

3 Preliminaries and Problem Formulation

This section describes flow matching as transport by a learned time-dependent vector field and formulates guided generation as control of that flow. It then states the path-constrained optimization problem and introduces HJB as the route to optimal control.

3 Preliminaries and Problem Formulation

  • Flow matching learns a time-dependent vector field that transports a noise distribution to a data distribution along trajectories defined by an ODE.Under regularity conditions, the vector field induces a flow map between states and a pushforward from the noise distribution to the data distribution.Source: §3 · pp. 3–4
  • Guided generation adds a chosen control input to a pretrained flow, then selects control to minimize a cost functional.The functional has running and terminal costs; one cited image-manipulation example uses a control penalty and negative CLIP similarity to encourage prompt adherence.Source: §3 · p. 4
  • The HJB equation characterizes the value function and yields an optimal control for the cost-minimization problem under its stated regularity conditions.The value function has a terminal condition given by the terminal loss.Source: §3 · pp. 4–5
  • The problem requires minimizing the cost while keeping the state within a time-varying safe set throughout the trajectory; myopic min-norm controllers can produce suboptimal solutions.Path-safety requires x(t) ∈ C_t for all t in the interval, rather than merely requiring the final sample to satisfy the constraint.Source: §3 · p. 5

4 Methodology

AegisFlow puts a time-varying barrier in a global cost to guide flow-matching trajectories safely over their full paths. It derives an HJB-based approximation to compute the resulting optimal control.

4.1 AegisFlow: Non-Myopic Path-Safe Guided Flow Matching

  • A time-varying barrier in the global cost penalizes approaching the safe-set boundary, and finite cost keeps the trajectory inside Ct; Proposition 2 formalizes safety under its stated conditions.The proposition assumes a smooth value function, a slowly varying Ct, and a barrier that grows sufficiently fast near the boundary.Source: §4.1 · pp. 5–6
  • The barrier weight βt controls the safety–editability trade-off, while λt controls the smoothness of the control input.Source: §4.1 · p. 5
  • With smooth V and λt > 0, Proposition 1 gives an explicit HJB equation, and the optimal control can be obtained from the value-function gradient ∇xV.Solving the HJB equation is generally intractable, motivating direct estimation of the control.Source: §4.1 · p. 6, §4.2 · p. 6
  • Path-safety of the exact optimal control does not automatically extend to the HJB-approximated control or Algorithm 1, so a safe guard such as MNC may be needed.Source: §4.1 · p. 6

4.2 Tractable HJB Approximation

  • A change of variables using the reference-flow map transforms the HJB equation into a forced Hamilton-Jacobi equation.Source: §4.2 · p. 6
  • For high-dimensional problems, the approximation uses Hopf-Lax when Γt and Gt are constant, then assumes a first-order terminal-loss expansion and Gt ≈ I; a unique minimizer yields the Moreau-Yosida gradient.Source: §4.2 · p. 7
  • The algorithm substitutes the iteratively updated controlled flow for the nominal flow, evaluating terminal loss along the controlled trajectory and reducing compounding linearization error.Source: §4.2 · p. 7

5 Experiments

Across image editing, Maze2D navigation, and robotic manipulation, AegisFlow is reported to improve target alignment or planning quality while maintaining safety. The experiments also examine its safety margins and behavior under novel constraints.

5.1 Text-Guided Image Manipulation with Identity Preservation

  • On CelebA-HQ, AegisFlow achieves the best CLIP and LPIPS scores among state-of-the-art methods, while MNC attains the highest ID.Figure 4 reports a larger safety margin that allows more aggressive target-prompt optimization.Source: §5.1 · p. 7, Table 1 · p. 8, Figure 4 · p. 8

5.2 Maze Navigation with Safety Constraints

  • On Maze2D, AegisFlow achieves the strongest safety margins and best generation quality, improving both safety metrics and reducing trajectory curvature and acceleration versus SafeFlowMatcher.Figure 5 shows a larger safety buffer early in the rollout.Source: §5.2 · p. 8, Table 2 · p. 9, Figure 5 · p. 9

5.3 Robotic Manipulation with Safety Constraints

  • On the Avoiding robotic-manipulation task, AegisFlow has the best goal success and safety rate while using the fewest timesteps.The task adds novel test-time constraints; Figure 6 visualizes trajectories avoiding forbidden regions.Source: §5.3 · pp. 9–10, Table 3 · p. 9, Figure 6 · p. 10

6 Conclusion

AegisFlow addresses myopic path-safe guidance through a global safety-critical control formulation and a tractable HJB approximation. The appendix background explains forward-invariant safe sets and the per-timestep min-norm controller approach.

6 Conclusion

  • AegisFlow integrates constraints into a global cost functional and derives a tractable HJB approximation to obtain optimal control.The conclusion reports empirical evaluations across diverse domains and identifies HJB approximation and computational efficiency as future-work areas.Source: §6 · p. 10

A.1 Control Barrier Functions

  • A time-varying safe set Ct is forward invariant when trajectories remain in it, with CBFs using the boundary behavior of h to prevent departure.The stated CBF result guarantees forward invariance for a Lipschitz controller in KCBF when the theorem's conditions hold.Source: §A.1 · p. 15

A.2 Min-norm Controller

  • The min-norm controller adds a safe control to nominal control and seeks to keep that correction as small as possible at each timestep.A slack variable relaxes the CBF constraint for stability and robustness.Source: §A.2 · p. 15

B.1 Proof of Proposition 1

The proof substitutes the cost functional and controlled dynamics into the HJB equation, then uses the quadratic dependence on control to obtain explicit value-function and control forms.

B.1 Proof of Proposition 1

  • Substituting the cost functional and dynamics into the HJB equation yields the expression used to derive the optimal control.The displayed derivation is partially fragmented in the supplied passage.Source: §B.1 · p. 16
  • The term minimized over the control is quadratic in u, allowing an explicit form for the value function and optimal control.The resulting HJB equation includes the drift, barrier cost, and a gradient-norm term.Source: §B.1 · p. 16
  • The proof concludes after deriving the HJB expression for the value function.Source: §B.1 · p. 16

B.2 Proof of Proposition 2

The proof establishes forward invariance by showing the value function cannot increase before a trajectory exits, while its barrier contribution would force it to diverge at the boundary. A separate derivation obtains the time derivative through the flow map.

B.2 Proof of Proposition 2

  • Along the optimal trajectory, the value function is non-increasing, contradicting its divergence at the boundary under the stated barrier and set-variation conditions.The derivative is bounded above by −β_tB, with β_t and B nonnegative before the first exit.Source: §B.2 · p. 16
  • Under the stated barrier lower bound and integrability conditions, the distance to the safe-set boundary cannot tend to zero at a finite first exit.The argument makes the reciprocal distance square-integrable, then uses absolute continuity to contradict the distance reaching zero.Source: §B.2 · p. 16

B.3 Proof of Proposition 3

  • The proof of Proposition 3 derives the value-function time derivative using the chain rule and the flow-map definition, then substitutes the identities into the governing equation.Source: §B.3 · p. 16

C Implementation Details

The implementation specifies task-specific safe sets, barrier functions, terminal losses, and schedules for image editing and two planning tasks. It also compares AegisFlow’s runtime with safe-guidance baselines.

C.1 Text-Guided Image Manipulation with Identity Preservation

  • Image editing defines safety through an ArcFace-distance safe set around the source image and uses a negative CLIP terminal loss to target the prompt.The setup uses a time-varying margin; λ_t is 0.6 and β_t is 0.1.Source: §C.1 · p. 17

C.2 Maze Navigation with Safety Constraints

  • Maze navigation trains a flow-matching prior and weights the BS2 barrier term more heavily because that constraint is harder to satisfy.Because environment interaction makes score non-differentiable, the terminal loss combines trajectory curvature and acceleration.Source: §C.2 · p. 17, Table 4 · p. 17

C.3 Robotic Manipulation with Safety Constraints

  • Robotic manipulation uses barriers for halfspace and circular constraints and a curvature terminal loss because environment-interaction timesteps are non-differentiable.The reported setup uses λ_t = 1.5 and β_t = 0.5.Source: §C.3 · pp. 17–18

D.1 Computational Efficiency Comparisons

  • AegisFlow takes slightly longer than OC-Flow but less time than OC-Flow + MNC for image manipulation, and is slower than SafeDiffuser and SafeFlowmatcher for maze navigation.The authors attribute the maze runtime difference to estimating the forced HJB solution rather than solving a quadratic program at each timestep.Source: §D.1 · p. 18

D.2 Metrics Comparisons for Text-Guided Image Manipulation

The section points readers to detailed text-guided image-manipulation metric comparisons on CelebA-HQ.

D.2 Metrics Comparisons for Text-Guided Image Manipulation

  • Detailed metric comparisons for text-guided image manipulation on CelebA-HQ are presented in Table 6.Source: §D.2 · p. 18

Related papers

Loading 2610.39825v1…