Source-linked AI summary
Trajectory Optimization via Schr\"{o}dinger Bridge Sampling
Mattia Mosso, Yang Liu, Heng Yang
TL;DR
The paper asks how trajectory optimization can retain distributional exploration while satisfying hard constraints exactly. It develops Euclidean and Riemannian ASBS formulations for soft- and hard-constrained problems, and reports effective experiments in both regimes.
Problem
The paper asks how to retain global, distributional exploration while satisfying hard trajectory constraints exactly, given that existing sampling and inference-based solvers generally use penalties without formally guaranteeing feasibility.
Source: §1 · pp. 1–2Method
The paper formulates KL-regularized optimization as Gibbs–Boltzmann sampling, using ASBS in rollout coordinates and Riemannian ASBS on hard-constraint feasibility manifolds.
Source: §1 · p. 2, §4.1 · p. 4, §4.3 · p. 5Results
Takeaways & Limitations
Hard-constrained trajectory optimization can be posed as sampling from an intrinsic Gibbs measure on the feasible manifold under the stated assumptions.
Source: §1 · p. 2, §4.2 · p. 4Computational complexity increases with the time horizon, or more precisely with the number of time steps T.
Source: §6 · p. 7
Abstract
from arXiv · showhide
We take a new look at the relation between finite-horizon trajectory optimization and Schr\"odinger bridge sampling. Viewed as inference, KL-regularized trajectory optimization is solved by sampling from a Gibbs--Boltzmann distribution whose energy is the trajectory cost, and the adjoint Schr\"odinger bridge sampler (ASBS) is a simulation-free diffusion sampler designed for exactly such unnormalized targets. Hard equality path and terminal constraints, by contrast, confine the admissible decision variables to a measure-zero feasibility manifold, on which the target must be redefined intrinsically. In particular: $(a)$ we analyze two complementary parametrizations; a rollout parametrization, in which the dynamics are eliminated and only the remaining constraints shape the manifold, and a double-shooting parametrization, in which states and controls are sampled jointly and the dynamics themselves become part of the manifold; $(b)$ we establish regularity conditions under which both admissible sets are smooth embedded manifolds; $(c)$ under compactness and path-connectedness assumptions, we sample from the resulting intrinsic Gibbs measures via Riemannian ASBS, treating strict inequalities through exponential slack variables. Experiments, including contact-rich locomotion and manipulation, demonstrate the effectiveness of both regimes.
1. INTRODUCTION
The introduction frames trajectory optimization as inference and asks whether sampling-based global exploration can be retained while enforcing hard constraints exactly. It presents ASBS for soft constraints and R–ASBS for hard-constrained feasibility manifolds.
1. INTRODUCTION
- Existing sampling and inference-based solvers generally encode constraints through soft penalties, trading feasibility against cost without formally guaranteeing feasibility.The paper asks whether global, distributional exploration can instead be retained under exact hard constraints.Source: §1 · pp. 1–2
- KL-regularized trajectory optimization optimizes a distribution over decision variables, yielding a Gibbs–Boltzmann target whose energy is induced by trajectory cost.As temperature λ approaches zero, the target concentrates on minimizers.Source: §1 · p. 2
- Soft-constrained rollout problems reduce to Euclidean Gibbs–Boltzmann sampling with ASBS, while hard-constrained problems become sampling on feasibility manifolds with R–ASBS.The hard-constrained target must be redefined intrinsically because the feasible set is measure zero in the ambient space.Source: §1 · p. 2
- For hard constraints, the paper establishes regularity conditions for smooth embedded feasible sets and assumes compactness and path-connectedness, with strict inequalities handled by exponential slacks.The stated conditions align with R–ASBS hypotheses.Source: §1 · p. 2
2. RELATED WORK
Related work spans optimization-based and sampling-based trajectory solvers, diffusion planning from demonstrations, and methods for sampling on constrained manifolds.
2. RELATED WORK
- Optimization-based methods include interior-point, sequential-quadratic, successive-convex, and augmented-Lagrangian programming, while mixed-integer and convex-relaxation alternatives offer certificates at rapidly growing computational cost.These methods appear in shooting or direct-transcription form.Source: §2 · p. 2
- Sampling-based trajectory optimization includes MPPI, predictive sampling, cross-entropy, Stein variational methods, model-based diffusion, and TSMC; constraints are generally encoded through costs, with constrained Stein variants enforcing them explicitly.TSMC’s inference formulation underlies the paper’s later development.Source: §2 · p. 2
- Diffusion models can synthesize trajectories, policies, or dynamics from demonstrations, with constraints imposed through guidance or projection and applicability resting on trajectory datasets.Source: §2 · p. 2
- Constrained Hamiltonian Monte Carlo is asymptotically exact for implicit manifolds but mixes slowly and requires a constraint solve at every chain step; mirror, reflected, and barrier methods also address constrained generative modeling.Source: §2 · p. 2
3. ADJOINT SCHRÖDINGER BRIDGE SAMPLING
This section presents Euclidean ASBS for sampling Gibbs–Boltzmann targets and its Riemannian extension for intrinsic sampling on manifolds. The manifold formulation adapts the sampler’s operations to geometric constraints, including implicitly defined manifolds.
3.1. Euclidean ASBS
- Euclidean ASBS seeks a controlled diffusion that transports a tractable source distribution to a target distribution, and the Schrödinger bridge is the minimum-energy control for that transport.The exposition restricts its reference process to zero drift.Source: §3.1 · p. 2
- ASBS alternates adjoint matching and corrector matching as simulation-free regressions, requiring neither differentiation through simulated trajectories nor samples from the target distribution.The corrector estimates endpoint bias from a non-memoryless source.Source: §3.1 · p. 3
3.2. Riemannian ASBS: Intrinsic Sampling on M
- Riemannian ASBS defines the target intrinsically on a smooth embedded manifold and evolves diffusion within its tangent bundle, preserving feasibility when initialized on the manifold.The formulation assumes a compact, path-connected manifold.Source: §3.2 · p. 3
- The general manifold sampler approximates heat-kernel scores with a short-time Varadhan/Log-map approximation, reference bridges with noisy geodesics, and parallel transport with Levi–Civita transport.Source: §3.2 · p. 3
- For implicitly defined nonlinear equality constraints, extended R–ASBS uses nearest-point projection for retraction, projection-as-transport for adjoints, and a projected-chord corrector target.The projection-based operations replace unavailable manifold primitives.Source: §3.2 · p. 3
4. TRAJECTORY OPTIMIZATION VIA SCHRÖDINGER BRIDGE SAMPLING
The section recasts soft-constrained trajectory optimization as Euclidean Gibbs–Boltzmann sampling and develops intrinsic sampling on feasibility manifolds for hard constraints. It presents rollout and double-shooting coordinates, regularity conditions for their feasible sets, and assumptions needed to apply R–ASBS.
4.1. Rollout Coordinates and the Soft-Constrained Case
- Soft-constrained optimization in rollout coordinates reduces to ASBS sampling from an unnormalized target, whose samples concentrate on low-cost controls as temperature falls.The dynamics enter through rollout derivatives of the energy.Source: §4.1 · p. 4
4.2. Hard Constraints
- Smooth dynamics and full row rank of the constraint Jacobian on a nonempty feasible set ensure the rollout feasible set is a smooth embedded manifold of dimension Tm − R.The rank condition is not automatic and couples stages through the rollout.Source: §4.2 · p. 4
- Because the rollout feasible set is measure zero in ambient control space, hard-constrained optimization must be posed intrinsically, with the uniform law as reference on a compact manifold.The constraints shape the manifold while dynamics are satisfied by construction.Source: §4.2 · p. 4
4.3. Double-Shooting Coordinates: Dynamics as Manifold Constraints
- Double shooting samples full trajectories with dynamics as explicit equality constraints; its dynamics manifold is smoothly embedded without a rank assumption, and the constrained manifold inherits rollout regularity.The constrained manifold is diffeomorphic to the rollout feasible set and has dimension Tm − R.Source: §4.3 · p. 5
- Double shooting makes constraint rows local in time and avoids exploding gradients, but increases ambient dimension and requires retraction onto a large nonlinear manifold at every step.Rollout coordinates instead use the smallest sampling dimension and preserve dynamic feasibility during projections, but derivatives can grow or vanish exponentially with horizon for unstable or stiff linearized dynamics.Source: §4.3 · p. 5, §4.2 · p. 5
4.4. Standing Assumptions and Strict Inequalities
- R–ASBS assumes a nonempty, compact, path-connected feasible set and access to a feasible source; strict inequalities are encoded by exponential slack variables.The slack construction represents only the strict interior, so boundary-active optima fall outside the targeted setting and can cause ill-conditioning.Source: §4.4 · pp. 5–6
5. NUMERICAL EXPERIMENTS
The experiments probe Euclidean ASBS on soft-constrained locomotion tasks, with differentiable rollout gradients and sampled trajectories. They cover planar half-cheetah and spatial quadruped benchmarks.
5. NUMERICAL EXPERIMENTS
- The experiments test the Euclidean and hard-constrained regimes beyond the theory section, using Python and JAX on a workstation with two NVIDIA Ada6000 GPUs.The source code is reported as available on GitHub.Source: §5 · p. 6
5.1. Soft-Constrained Tasks via Euclidean ASBS
- Reverse-mode differentiation through the contact rollout produces ∇zJ, agreeing with central finite differences across four decades of perturbation size.The evaluation reports 2,000 sampled trajectories per evaluation and selects the lowest-cost trajectory.Source: §5.1 · p. 6
5.1.1. Cheetah Locomotion:
- The planar half-cheetah uses 150 control knots held for 10 integration steps, yielding a sampled variable z in R^900 and 1,500 contact steps per rollout.Its running cost combines deviation from the benchmark forward-speed target with small control-effort and smoothness terms.Source: §5.1.1 · p. 6
5.1.2. Quadruped Locomotion:
- The spatial quadruped has 12 actuators across four alternating stance-and-flight legs, with 150 knots held for four integration steps and a sampled variable z in R^1800.The running cost is the benchmark Move objective.Source: §5.1.2 · p. 6
5.2. Hard Equality Constraints via R–ASBS
The hard-equality experiments apply R–ASBS to navigation, cooperative transport, and object pushing. They specify task constraints and report selected low-cost sampled trajectories, without supplied quantitative comparisons of outcomes.
5.2.1. 2D PointMaze:
- PointMaze samples trajectories for a point mass with a fixed start and exact terminal arrival; it compares against constrained NLP and RGD from the same prior in double-shooting coordinates.The maze is encoded through the running cost, and Fig. 2 marks the initial state and hard terminal goal.Source: §5.2.1 · p. 6, Figure 2 · p. 7
5.2.2. Multi-Drone Cooperative Transport:
- Three UAVs carry a cable-suspended payload through two timed waypoints while avoiding moving obstacles and inter-drone collisions; cable angles represent fixed-length geometry.The full initial and terminal state and waypoint passages are hard equalities; other specifications are soft running costs, and the selected trajectory is chosen from 2,500 samples.Source: §5.2.2 · pp. 6–7
5.2.3. Push-T Task:
- Push-T models a prescribed-contact pusher and T-shaped slider with a soft Gumbel–Softmax segment selector; target pose is enforced by three hard terminal equalities.Control smoothness and obstacle avoidance are soft costs, and the figure reports the lowest-cost trajectory among 500 samples.Source: §5.2.3 · p. 7, Figure 3 · p. 7
5.3. Strict Inequality Constraints
The AUV experiment evaluates the slack construction on a constrained subsea pipeline-inspection task, reporting a selected low-cost trajectory from sampled candidates.
5.3. Strict Inequality Constraints
- The AUV must reach a terminal target through an inspection gate while staying in a safe pipeline annulus, avoiding adverse flow and divers, and remaining above the seafloor.The task also specifies safety distance from two divers working near the pipe.Source: §5.3 · p. 7
- The rollout uses bounded raw-acceleration controls and includes position- and time-dependent ocean current, nonlinear hydrodynamic drag, and a small vertical bias.Physical acceleration is bounded through a tanh transformation of the raw control.Source: §5.3 · p. 7
- For fixed work zones centered at ok, k = 1, 2, Fig. 4(a) reports the selected lowest-cost trajectory out of 2,500 sampled ones.Source: §5.3 · p. 7
6. CONCLUSION
The conclusion presents the work as a connection between finite-horizon trajectory optimization and diffusion Schrödinger bridge sampling, with a sampling framework for hard-constrained instances. It reports experimental efficacy and identifies time-step growth as a limitation.
6. CONCLUSION
- The authors report that their sampling-based framework for hard-constrained instances was effective in numerical experiments.The conclusion does not provide a quantitative measure in this passage.Source: §6 · p. 7
- Computational complexity increases with the time horizon, more precisely with the number of time steps T.The authors leave this limitation for future work.Source: §6 · p. 7
- Fig. 4 depicts AUV pipeline inspection and multi-drone payload transport in front, top, and side views, with feasible trajectory samples and dashed feasible-source means.The caption identifies the two scenarios as panels (a) and (b), respectively.Source: Figure 4 · p. 8
APPENDIX A SUPPLEMENTARY MATERIALS
The appendix establishes the geometry of the rollout and constrained feasible sets, then shows that the intrinsic Gibbs measure is the unique solution of the corresponding KL-regularized problem. It also gives stagewise sufficient conditions for constraint regularity.
Intrinsic Gibbs measure
- On compact M, continuity of E bounds the normalizing constant above and below by positive finite values.The Gibbs density is bounded and strictly positive with respect to the reference measure.Source: Appendix A · p. 8
- The intrinsic Gibbs measure uniquely minimizes the KL-regularized objective on the manifold, with minimum value −λ log(ZM/VM).The proof rewrites the objective as a KL divergence to that measure plus a constant.Source: Appendix A · p. 8
Rollout geometry
- The rollout map is a diffeomorphism onto the dynamics-feasible set, which is an embedded manifold of dimension n + Tm without a dynamics-rank condition.Its derivative is injective because the projection back to initial state and controls is a left inverse.Source: Appendix A · p. 9
Constraint regularity
- With a fixed initial state, the path- and terminal-constrained rollout set is an embedded manifold of dimension Tm − Rp − rT when the stated regularity conditions hold.The stagewise sufficient condition checks full row rank of each control-constraint Jacobian and surjectivity of the terminal map on the path-constrained tangent space.Source: Appendix A · pp. 9–10
- The path-constraint Jacobian has block lower-triangular structure because each stage constraint is independent of its current control through the preceding state.Full row rank of each diagonal block yields the required rank for the path-constraint map.Source: Appendix A · p. 9
APPENDIX B ALGORITHM STRUCTURES
The appendix gives algorithm structures for Euclidean ASBS and its manifold extension, including alternating regressions, projected-and-retracted feasible sources, and implementation choices for initialization and sampling.
Algorithm 1
- Algorithm 1 alternates two regressions, reusing simulated endpoints within each epoch so energy-gradient evaluations depend on buffer size rather than parameter-update count.The endpoint buffer is populated once per epoch.Source: Appendix B · p. 10
Implementation details
- The adjoint corrector starts at zero, and terminal-gradient norm clipping bounds regression targets for heavy-tailed gradients without changing the stationary target.Buffers are populated without differentiating through the sampling SDE.Source: Appendix B · pp. 10–11
Algorithm 2
- The manifold sampler's feasible source perturbs an approximately feasible nominal point in its tangent space and retracts the perturbation onto the manifold.The width κ controls the residual that Newton iterations must absorb; failed retractions are monitored rather than rejected.Source: Appendix B · pp. 10–11
- In rollout coordinates, the nominal point can be built by interpolating a reference path, inverting the dynamics for controls, and closing inequality slacks in closed form.The same nominal point and covariance define the Gaussian energy term and a prior centered on the nominal trajectory.Source: Appendix B · p. 11