diff --git a/Project.toml b/Project.toml index 33ee19e2..e4d9552b 100644 --- a/Project.toml +++ b/Project.toml @@ -19,12 +19,12 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe" TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8" [compat] -DirectTrajOpt = "0.5" +DirectTrajOpt = "0.6" Distributions = "0.25" ExponentialAction = "0.2" LinearAlgebra = "1.10, 1.11, 1.12" -NamedTrajectories = "0.6" -PiccoloQuantumObjects = "0.7" +NamedTrajectories = "0.7" +PiccoloQuantumObjects = "0.8" Random = "1.10, 1.11, 1.12" Reexport = "1.2" SparseArrays = "1.10, 1.11, 1.12" diff --git a/docs/literate/examples/minimum_time_problem.jl b/docs/literate/examples/minimum_time_problem.jl new file mode 100644 index 00000000..e0705c3b --- /dev/null +++ b/docs/literate/examples/minimum_time_problem.jl @@ -0,0 +1,157 @@ +# # Minimum Time Optimization + +# In this example, we'll show how to optimize for the **fastest possible gate** while maintaining high fidelity. This is useful when you want to minimize gate duration to reduce decoherence effects. + +# ## Problem Setup + +# We'll start with a smooth pulse solution for an X gate, then optimize for minimum time. + +using QuantumCollocation +using PiccoloQuantumObjects +using NamedTrajectories + +using PiccoloPlots +using CairoMakie + +# Define a simple qubit system: + +H_drift = 0.05 * PAULIS.Z +H_drives = [PAULIS.X, PAULIS.Y] +T_max = 50.0 # Start with generous time budget +drive_bounds = [(-0.5, 0.5), (-0.5, 0.5)] + +sys = QuantumSystem(H_drift, H_drives, T_max, drive_bounds) + +# Target gate: X gate (π rotation about X axis) + +U_goal = GATES.X +N = 101 # Use more timesteps for better time resolution + +# ## Step 1: Find a Smooth Pulse Solution + +# First, we solve for a smooth pulse with a fixed time to get a good starting point: + +prob_smooth = UnitarySmoothPulseProblem(sys, U_goal, N) +solve!(prob_smooth; max_iter=100) + +fid_smooth = unitary_rollout_fidelity(prob_smooth.trajectory, sys) +duration_smooth = get_duration(prob_smooth.trajectory) + +println("Smooth pulse fidelity: ", fid_smooth) +println("Smooth pulse duration: ", duration_smooth) + +# Let's visualize the smooth pulse solution: + +plot_unitary_populations(prob_smooth.trajectory) + +# ## Step 2: Optimize for Minimum Time + +# Now we'll use the smooth pulse as a starting point and optimize for minimum time while constraining the fidelity: + +prob_mintime = UnitaryMinimumTimeProblem( + prob_smooth, # Use smooth pulse as initial guess + U_goal; + final_fidelity=0.9999 # Require high fidelity +) + +# Solve the minimum time problem: + +solve!(prob_mintime; max_iter=300) + +# ## Results + +# Let's compare the results: + +fid_mintime = unitary_rollout_fidelity(prob_mintime.trajectory, sys) +duration_mintime = get_duration(prob_mintime.trajectory) + +println("\n=== Comparison ===") +println("Smooth pulse:") +println(" Duration: ", duration_smooth) +println(" Fidelity: ", fid_smooth) +println("\nMinimum time:") +println(" Duration: ", duration_mintime) +println(" Fidelity: ", fid_mintime) +println("\nSpeedup: ", duration_smooth / duration_mintime, "×") + +@assert fid_mintime > 0.9999 +@assert duration_mintime < duration_smooth + +# Visualize the minimum time solution: + +plot_unitary_populations(prob_mintime.trajectory) + +# ## Understanding the Tradeoff + +# Let's look at how the control pulses changed: + +fig = Figure(size=(1000, 600)) + +## Smooth pulse controls +ax1 = Axis(fig[1, 1], xlabel="Time", ylabel="Control amplitude", title="Smooth Pulse") +for i in 1:size(prob_smooth.trajectory.u, 1) + lines!(ax1, prob_smooth.trajectory.u[i, :], label="u$i") +end +axislegend(ax1) + +## Minimum time controls +ax2 = Axis(fig[1, 2], xlabel="Time", ylabel="Control amplitude", title="Minimum Time") +for i in 1:size(prob_mintime.trajectory.u, 1) + lines!(ax2, prob_mintime.trajectory.u[i, :], label="u$i") +end +axislegend(ax2) + +fig + +# Notice how the minimum time solution: +# - Uses fewer timesteps (compressed duration) +# - Has more aggressive control amplitudes +# - May have sharper transitions + +# ## Varying the Fidelity Constraint + +# Let's see how gate duration depends on the required fidelity: + +fidelities = [0.99, 0.995, 0.999, 0.9999] +durations = Float64[] + +for target_fid in fidelities + prob = UnitaryMinimumTimeProblem( + prob_smooth, + U_goal; + final_fidelity=target_fid + ) + solve!(prob; max_iter=200, verbose=false) + push!(durations, get_duration(prob.trajectory)) + println("Fidelity ", target_fid, " → Duration ", durations[end]) +end + +# Plot the tradeoff: + +fig_tradeoff = Figure() +ax = Axis( + fig_tradeoff[1, 1], + xlabel="Required Fidelity", + ylabel="Gate Duration", + title="Fidelity vs Duration Tradeoff" +) +scatter!(ax, fidelities, durations, markersize=15) +lines!(ax, fidelities, durations) +fig_tradeoff + +# ## Key Takeaways + +# 1. **Start with smooth pulse** - Good initialization is crucial +# 2. **Higher fidelity = longer duration** - There's always a tradeoff +# 3. **More timesteps = better resolution** - Use enough N for time optimization +# 4. **Control bounds matter** - Tighter bounds → longer minimum time +# +# ## When to Use Minimum Time Problems +# +# Use `UnitaryMinimumTimeProblem` when: +# - Decoherence is limiting your gate fidelity +# - You want to maximize throughput in quantum circuits +# - You need to compare against theoretical speed limits +# - You're exploring the fundamental limits of your control system +# +# For most applications, `UnitarySmoothPulseProblem` with reasonable duration is sufficient and more robust. diff --git a/docs/literate/examples/multilevel_transmon.jl b/docs/literate/examples/multilevel_transmon.jl index 5f3d02a7..6b22f504 100644 --- a/docs/literate/examples/multilevel_transmon.jl +++ b/docs/literate/examples/multilevel_transmon.jl @@ -35,22 +35,18 @@ using CairoMakie ## define the time parameters -T₀ = 10.0 # total time in ns N = 50 # number of time steps -Δt = T₀ / N # time step +T₀ = 10.0 # total time in ns ## define the system parameters levels = 5 δ = 0.2 ## add a bound to the controls -u_bounds = [0.2, 0.2] +u_bound = 0.2 ## create the system -sys = TransmonSystem(levels=levels, δ=δ, T_max=T₀, u_bounds=u_bounds) - -## let's look at a drive -get_drives(sys)[1] |> sparse +sys = TransmonSystem(levels=levels, δ=δ, T_max=T₀, drive_bounds=fill(u_bound, 2)) # Since this is a multilevel transmon and we want to implement an, let's say, $X$ gate on the qubit subspace, i.e., the first two levels we can utilize the `EmbeddedOperator` type to define the target operator. @@ -75,16 +71,16 @@ get_subspace_identity(op) |> sparse # We can then pass this embedded operator to the `UnitarySmoothPulseProblem` template to create the problem ## create the problem -prob = UnitarySmoothPulseProblem(sys, op, N, Δt) +prob = UnitarySmoothPulseProblem(sys, op, N) ## solve the problem -load_path = joinpath(dirname(Base.active_project()), "data/multilevel_transmon_example_89ee72.jld2") # hide -prob.trajectory = load_traj(load_path) # hide -nothing # hide +# load_path = joinpath(dirname(Base.active_project()), "data/multilevel_transmon_example_89ee72.jld2") # hide +# prob.trajectory = load_traj(load_path) # hide +# nothing # hide +# solve!(prob; max_iter=50) #= ```julia -solve!(prob; max_iter=50) ``` ```@raw html @@ -179,7 +175,7 @@ plot_unitary_populations(prob.trajectory; fig_size=(900, 700)) ## create the a leakage suppression problem, initializing with the previous solution -prob_leakage = UnitarySmoothPulseProblem(sys, op, N, Δt; +prob_leakage = UnitarySmoothPulseProblem(sys, op, N; u_guess=prob.trajectory.u[:, :], piccolo_options=PiccoloOptions( leakage_constraint=true, @@ -189,13 +185,12 @@ prob_leakage = UnitarySmoothPulseProblem(sys, op, N, Δt; ) ## solve the problem -load_path = joinpath(dirname(Base.active_project()), "data/multilevel_transmon_example_leakage_89ee72.jld2") # hide -prob_leakage.trajectory = load_traj(load_path) # hide -nothing # hide - +# load_path = joinpath(dirname(Base.active_project()), "data/multilevel_transmon_example_leakage_89ee72.jld2") # hide +# prob_leakage.trajectory = load_traj(load_path) # hide +# nothing # hide +# solve!(prob_leakage; max_iter=250, options=IpoptOptions(eval_hessian=false)) #= ```julia -solve!(prob_leakage; max_iter=250) ``` ```@raw html diff --git a/docs/literate/examples/robust_control.jl b/docs/literate/examples/robust_control.jl new file mode 100644 index 00000000..a5eee4f1 --- /dev/null +++ b/docs/literate/examples/robust_control.jl @@ -0,0 +1,264 @@ +# # Robust Control with Sampling + +# Real quantum systems have uncertainties - manufacturing variations, calibration errors, environmental fluctuations. **Robust control** finds pulses that work well across these variations. + +# In this example, we'll use `UnitarySamplingProblem` to optimize a gate that's robust to: +# - Drift Hamiltonian uncertainties +# - Control amplitude miscalibrations + +# ## Problem Setup + +using QuantumCollocation +using PiccoloQuantumObjects +using NamedTrajectories + +using PiccoloPlots +using CairoMakie + +# ## Nominal System + +# First, define our nominal (ideal) qubit system: + +H_drift_nominal = 0.1 * PAULIS.Z +H_drives = [PAULIS.X, PAULIS.Y] +T_max = 15.0 +drive_bounds = [(-1.0, 1.0), (-1.0, 1.0)] + +sys_nominal = QuantumSystem(H_drift_nominal, H_drives, T_max, drive_bounds) + +# Target: Hadamard gate + +U_goal = GATES.H +N = 51 + +# ## Non-Robust Solution (Baseline) + +# Let's first solve without robustness to see how a standard solution performs: + +prob_standard = UnitarySmoothPulseProblem(sys_nominal, U_goal, N) +solve!(prob_standard; max_iter=100) + +fid_standard = unitary_rollout_fidelity(prob_standard.trajectory, sys_nominal) +println("Standard solution fidelity (nominal system): ", fid_standard) + +# ## System Variations + +# Now let's model realistic system variations. We'll consider: +# 1. **Drift variation**: ±10% uncertainty in the Z bias +# 2. **Control miscalibration**: ±5% error in drive amplitudes + +# Create systems representing these variations: + +## Drift variations +drift_variations = [-0.1, 0.0, 0.1] # ±10% +systems = QuantumSystem[] + +for δ in drift_variations + H_drift_varied = (1.0 + δ) * H_drift_nominal + sys = QuantumSystem(H_drift_varied, H_drives, T_max, drive_bounds) + push!(systems, sys) +end + +println("Number of system variations: ", length(systems)) + +# ## Robust Optimization + +# Now use `UnitarySamplingProblem` to optimize over all system variations simultaneously: + +## Create target operators for each system (all want the same gate) +operators = [EmbeddedOperator(U_goal, sys) for sys in systems] + +## Set up robust problem +prob_robust = UnitarySamplingProblem( + systems, + operators, + N; + system_weights=fill(1.0, length(systems)) # Equal weight to all variations +) + +# Solve the robust problem: + +solve!(prob_robust; max_iter=200) + +# ## Compare Performance + +# Let's evaluate both solutions (standard and robust) across all system variations: + +println("\n=== Fidelity Comparison ===") +println("System Variation | Standard | Robust") +println("-" ^ 50) + +fidelities_standard = Float64[] +fidelities_robust = Float64[] + +for (i, sys) in enumerate(systems) + ## Standard solution on this system + traj_std = prob_standard.trajectory + fid_std = unitary_rollout_fidelity(traj_std, sys) + + ## Robust solution on this system + traj_rob = prob_robust.trajectory + fid_rob = unitary_rollout_fidelity(traj_rob, sys) + + push!(fidelities_standard, fid_std) + push!(fidelities_robust, fid_rob) + + println("Drift $(drift_variations[i]): | $(round(fid_std, digits=5)) | $(round(fid_rob, digits=5))") +end + +println("\nWorst-case fidelity:") +println(" Standard: ", minimum(fidelities_standard)) +println(" Robust: ", minimum(fidelities_robust)) + +# The robust solution should have better worst-case fidelity! + +# ## Visualize the Tradeoff + +# Plot fidelities across variations: + +fig = Figure(size=(800, 500)) +ax = Axis( + fig[1, 1], + xlabel="Drift Variation (%)", + ylabel="Gate Fidelity", + title="Robustness to System Variations" +) + +x_vals = drift_variations .* 100 + +scatter!(ax, x_vals, fidelities_standard, markersize=15, label="Standard", color=:red) +lines!(ax, x_vals, fidelities_standard, color=:red, linestyle=:dash) + +scatter!(ax, x_vals, fidelities_robust, markersize=15, label="Robust", color=:blue) +lines!(ax, x_vals, fidelities_robust, color=:blue) + +axislegend(ax) +fig + +# ## Control Pulse Comparison + +# How do the control pulses differ? + +fig_controls = Figure(size=(1000, 600)) + +## Standard solution +ax1 = Axis(fig_controls[1, 1], xlabel="Timestep", ylabel="Amplitude", title="Standard (Non-Robust)") +for i in 1:size(prob_standard.trajectory.u, 1) + lines!(ax1, prob_standard.trajectory.u[i, :], label="u$i") +end +axislegend(ax1) + +## Robust solution +ax2 = Axis(fig_controls[1, 2], xlabel="Timestep", ylabel="Amplitude", title="Robust (Sampling)") +for i in 1:size(prob_robust.trajectory.u, 1) + lines!(ax2, prob_robust.trajectory.u[i, :], label="u$i") +end +axislegend(ax2) + +fig_controls + +# Robust pulses often: +# - Have slightly lower peak amplitudes +# - Are more "conservative" to handle uncertainties +# - May use different pulse shapes + +# ## More Realistic Variations + +# Let's test with more system variations including control miscalibration: + +## Create a grid of variations +drift_vals = [-0.15, -0.05, 0.0, 0.05, 0.15] +control_scale_vals = [0.95, 1.0, 1.05] + +systems_extended = QuantumSystem[] +operators_extended = EmbeddedOperator[] + +for drift_δ in drift_vals + for ctrl_scale in control_scale_vals + H_drift = (1.0 + drift_δ) * H_drift_nominal + H_drives_scaled = [ctrl_scale * H for H in H_drives] + sys = QuantumSystem(H_drift, H_drives_scaled, T_max, drive_bounds) + push!(systems_extended, sys) + push!(operators_extended, EmbeddedOperator(U_goal, sys)) + end +end + +println("\nExtended variations: ", length(systems_extended), " systems") + +# Optimize with all variations: + +prob_robust_extended = UnitarySamplingProblem( + systems_extended, + operators_extended, + N; + system_weights=fill(1.0, length(systems_extended)) +) + +solve!(prob_robust_extended; max_iter=300) + +# Evaluate performance: + +fidelities_extended = [ + unitary_rollout_fidelity(prob_robust_extended.trajectory, sys) + for sys in systems_extended +] + +println("\nExtended robust optimization:") +println(" Mean fidelity: ", mean(fidelities_extended)) +println(" Worst-case fidelity: ", minimum(fidelities_extended)) +println(" Best-case fidelity: ", maximum(fidelities_extended)) + +# ## Weighted Sampling + +# You can also prioritize certain variations by adjusting `system_weights`: + +## Give more weight to nominal system +weights_prioritized = ones(length(systems)) +weights_prioritized[2] = 5.0 # Nominal system (middle of variation range) + +prob_weighted = UnitarySamplingProblem( + systems, + operators, + N; + system_weights=weights_prioritized +) + +solve!(prob_weighted; max_iter=200) + +# Compare: + +fidelities_weighted = [ + unitary_rollout_fidelity(prob_weighted.trajectory, sys) + for sys in systems +] + +println("\n=== Weighting Comparison ===") +println("System | Equal Weight | Prioritized") +println("-" ^ 45) +for i in 1:length(systems) + println("$(i) | $(round(fidelities_robust[i], digits=5)) | $(round(fidelities_weighted[i], digits=5))") +end + +# ## Key Takeaways + +# 1. **Use UnitarySamplingProblem** for robust control +# 2. **Model realistic variations** - drift, control scaling, decoherence +# 3. **Balance fidelity vs robustness** - more variations = more conservative pulses +# 4. **Weight critical scenarios** - prioritize what matters most +# 5. **Test on unseen variations** - verify robustness generalizes +# +# ## When to Use Robust Control +# +# Use sampling-based robust control when: +# - System parameters have significant uncertainties +# - Gates must work across many devices (fab variations) +# - Environmental conditions fluctuate +# - Calibration drifts between experiments +# - Deploying on hardware without frequent recalibration +# +# ## Alternatives +# +# For other robustness approaches, see: +# - `UnitaryVariationalProblem` - Variational robustness with Hamiltonian perturbations +# - `UnitaryFreePhaseProblem` - Robustness to global phase errors +# - Custom constraints for specific robustness requirements diff --git a/docs/literate/examples/single_qubit_gate.jl b/docs/literate/examples/single_qubit_gate.jl new file mode 100644 index 00000000..dac0079c --- /dev/null +++ b/docs/literate/examples/single_qubit_gate.jl @@ -0,0 +1,101 @@ +# # Single Qubit Gate + +# This is the canonical example for getting started with QuantumCollocation.jl. We'll optimize control pulses to implement a Hadamard gate on a single qubit. + +# ## Problem Setup + +# We have a qubit with a drift Hamiltonian and two control drives (X and Y). Our goal is to find smooth control pulses that implement the Hadamard gate. + +using QuantumCollocation +using PiccoloQuantumObjects +using NamedTrajectories + +using PiccoloPlots +using CairoMakie + +# Define the system Hamiltonian: $H(t) = H_{\text{drift}} + u_x(t) H_x + u_y(t) H_y$ + +## Drift (always on, not controllable) +H_drift = 0.1 * PAULIS.Z # Small Z bias + +## Control Hamiltonians (we can modulate these) +H_drives = [PAULIS.X, PAULIS.Y] + +## System parameters +T_max = 10.0 # Maximum time (ns or μs, depends on your units) +drive_bounds = [(-1.0, 1.0), (-1.0, 1.0)] # Control amplitude bounds + +## Create the quantum system +sys = QuantumSystem(H_drift, H_drives, T_max, drive_bounds) + +# ## Define the Target Gate + +# We want to implement a Hadamard gate: + +U_goal = GATES.H +println("Target gate:") +display(U_goal) + +# ## Set Up the Optimization Problem + +# We'll use 51 timesteps to discretize the time interval: + +N = 51 # Number of timesteps + +## Create a smooth pulse problem +prob = UnitarySmoothPulseProblem(sys, U_goal, N) + +# Let's check the initial fidelity before optimization: + +fid_initial = unitary_rollout_fidelity(prob.trajectory, sys) +println("Initial fidelity: ", fid_initial) + +# ## Solve the Problem + +solve!(prob; max_iter=100) + +# ## Analyze the Results + +# Check the final fidelity: + +fid_final = unitary_rollout_fidelity(prob.trajectory, sys) +println("Final fidelity: ", fid_final) +@assert fid_final > 0.99 + +# Plot the optimized control pulses: + +plot_unitary_populations(prob.trajectory) + +# ## Understanding the Solution + +# Let's look at the control pulses in more detail: + +## Extract the controls +u = prob.trajectory.u # Control amplitudes (2 × N matrix) +du = prob.trajectory.du # First derivatives +ddu = prob.trajectory.ddu # Second derivatives + +println("Control pulse shape: ", size(u)) +println("Max control amplitude: ", maximum(abs.(u))) +println("Max control derivative: ", maximum(abs.(ddu))) + +# The smooth pulse problem penalizes both control amplitude and its derivatives, resulting in smooth, implementable pulses. + +# ## Duration + +# What's the total duration of this gate? + +duration = get_duration(prob.trajectory) +println("Gate duration: ", duration) + +# ## Next Steps + +# Now that you've seen the basic workflow, you can: +# +# - Try different gates (X, Y, Z, CNOT on coupled qubits) +# - Adjust the number of timesteps N +# - Modify the regularization weights (R_u, R_du, R_ddu) +# - Add constraints like leakage suppression for multilevel systems +# - Optimize for minimum time instead of smooth pulses +# +# See the other examples and the problem templates documentation for more! diff --git a/docs/literate/examples/two_qubit_gates.jl b/docs/literate/examples/two_qubit_gates.jl index c6f15aec..f60a299e 100644 --- a/docs/literate/examples/two_qubit_gates.jl +++ b/docs/literate/examples/two_qubit_gates.jl @@ -13,7 +13,7 @@ # Specifically, for a simple two-qubit system in a rotating frame, we have # ```math -# H = J_{12} \sigma_1^x \sigma_2^x + \sum_{i \in {1,2}} a_i^R(t) {\sigma^x_i \over 2} + a_i^I(t) {\sigma^y_i \over 2}. +# H = J_{12} \sigma_1^x \sigma_2^x + \sum_{i \in {1,2}} u_i^R(t) {\sigma^x_i \over 2} + u_i^I(t) {\sigma^y_i \over 2}. # ``` # where @@ -21,7 +21,7 @@ # ```math # \begin{align*} # J_{12} &= 0.001 \text{ GHz}, \\ -# |a_i^R(t)| &\leq 0.1 \text{ GHz} \\ +# |u_i^R(t)| &\leq 0.1 \text{ GHz} \\ # \end{align*} # ``` @@ -40,9 +40,9 @@ using CairoMakie ⊗(a, b) = kron(a, b) ## Define our operators -σx = GATES[:X] -σy = GATES[:Y] -Id = GATES[:I] +σx = GATES.X +σy = GATES.Y +Id = GATES.I ## Lift the operators to the two-qubit Hilbert space σx_1 = σx ⊗ Id @@ -61,7 +61,8 @@ H_drift = J_12 * (σx ⊗ σx) ## Define the control Hamiltonians H_drives = [σx_1 / 2, σy_1 / 2, σx_2 / 2, σy_2 / 2] -## Define higher derivative bounds +## Define control (and higher derivative) bounds +u_bound = 0.1 du_bound = 0.0005 ddu_bound = 0.0025 @@ -71,14 +72,11 @@ H_drives .*= 2π ## Define the time parameters N = 100 # timesteps -T_max = 1.0 # max evolution time -u_bounds = fill((-u_bound, u_bound), length(H_drives)) -duration = 100 # μs -Δt = duration / N -Δt_max = 400 / N +T_max = 400.0 # μs (maximum duration) +drive_bounds = fill((-u_bound, u_bound), length(H_drives)) ## Define the system -sys = QuantumSystem(H_drift, H_drives, T_max, u_bounds) +sys = QuantumSystem(H_drift, H_drives, T_max, drive_bounds) # ## SWAP gate @@ -93,27 +91,25 @@ U_goal = [ ## Set up the problem prob = UnitarySmoothPulseProblem( sys, - U_goal, - N, - Δt; + EmbeddedOperator(U_goal, sys), + N; du_bound=du_bound, ddu_bound=ddu_bound, R_du=0.01, R_ddu=0.01, - Δt_max=Δt_max, piccolo_options=PiccoloOptions(bound_state=true), ) fid_init = unitary_rollout_fidelity(prob.trajectory, sys) println(fid_init) # Solve the problem -load_path = joinpath(dirname(Base.active_project()), "data/two_qubit_gates_89ee72.jld2") # hide -prob.trajectory = load_traj(load_path) # hide -nothing # hide +# load_path = joinpath(dirname(Base.active_project()), "data/two_qubit_gates_89ee72.jld2") # hide +# prob.trajectory = load_traj(load_path) # hide +# nothing # hide +# solve!(prob; max_iter=100, options=IpoptOptions(eval_hessian=false)) #= ```julia -solve!(prob; max_iter=100) ``` ```@raw html @@ -207,13 +203,13 @@ plot_unitary_populations(prob.trajectory) min_time_prob = UnitaryMinimumTimeProblem(prob, U_goal; final_fidelity=.995) # and solve the problem -load_path = joinpath(dirname(Base.active_project()), "data/two_qubit_gates_min_time_89ee72.jld2") # hide -min_time_prob.trajectory = load_traj(load_path) # hide -nothing # hide +# load_path = joinpath(dirname(Base.active_project()), "data/two_qubit_gates_min_time_89ee72.jld2") # hide +# min_time_prob.trajectory = load_traj(load_path) # hide +# nothing # hide +# solve!(min_time_prob; max_iter=300) #= ```julia -solve!(min_time_prob; max_iter=300) ``` ```@raw html
    initializing optimizer...
@@ -322,27 +318,25 @@ U_goal_ms = exp(im * π/4 * σx_1 * σx_2)
 
 prob_ms = UnitarySmoothPulseProblem(
     sys,
-    U_goal_ms,
-    N,
-    Δt;
+    EmbeddedOperator(U_goal, sys),
+    N;
     du_bound=du_bound,
     ddu_bound=ddu_bound,
     R_du=0.01,
     R_ddu=0.01,
-    Δt_max=Δt_max,
     piccolo_options=PiccoloOptions(bound_state=true),
 )
 fid_init = unitary_rollout_fidelity(prob_ms.trajectory, sys)
 println(fid_init)
 
 # Solve the problem
-load_path = joinpath(dirname(Base.active_project()), "data/two_qubit_gates_molmer_89ee72.jld2") # hide
-prob_ms.trajectory = load_traj(load_path) # hide
-nothing # hide
+# load_path = joinpath(dirname(Base.active_project()), "data/two_qubit_gates_molmer_89ee72.jld2") # hide
+# prob_ms.trajectory = load_traj(load_path) # hide
+# nothing # hide
+# solve!(prob; max_iter=300)
 
 #=
 ```julia
-solve!(prob; max_iter=300)
 ```
 
 ```@raw html
    initializing optimizer...
@@ -427,13 +421,13 @@ plot_unitary_populations(prob_ms.trajectory)
 min_time_prob_ms = UnitaryMinimumTimeProblem(prob_ms, U_goal_ms; final_fidelity=.9995)
 
 # and solve the problem
-load_path = joinpath(dirname(Base.active_project()), "data/two_qubit_gates_molmer_min_time_89ee72.jld2") # hide
-min_time_prob_ms.trajectory = load_traj(load_path) # hide
-nothing # hide
+# load_path = joinpath(dirname(Base.active_project()), "data/two_qubit_gates_molmer_min_time_89ee72.jld2") # hide
+# min_time_prob_ms.trajectory = load_traj(load_path) # hide
+# nothing # hide
+# solve!(min_time_prob; max_iter=300)
 
 #=
 ```julia
-solve!(min_time_prob; max_iter=300)
 ```
 
 ```@raw html
   initializing optimizer...
diff --git a/docs/literate/man/adding_constraints.jl b/docs/literate/man/adding_constraints.jl
new file mode 100644
index 00000000..23fe16d7
--- /dev/null
+++ b/docs/literate/man/adding_constraints.jl
@@ -0,0 +1,310 @@
+# # Adding Custom Constraints
+
+# QuantumCollocation.jl provides problem templates with built-in constraints, but you often need to add custom constraints for specific physics or hardware requirements.
+
+# This guide shows you how to add:
+# - Path constraints (applied at every timestep)
+# - Boundary constraints (initial/final conditions)
+# - Nonlinear constraints (complex physics)
+# - Leakage suppression
+# - Control envelope constraints
+
+using QuantumCollocation
+using PiccoloQuantumObjects
+using NamedTrajectories
+using DirectTrajOpt
+
+# ## Basic Setup
+
+# Start with a simple system:
+
+H_drift = 0.1 * PAULIS.Z
+H_drives = [PAULIS.X, PAULIS.Y]
+sys = QuantumSystem(H_drift, H_drives, 10.0, [1.0, 1.0])
+U_goal = GATES.H
+N = 51
+
+# ## 1. Path Constraints (Box Constraints)
+
+# The simplest constraints are bounds on variables at all timesteps.
+
+# ### Built-in Bounds
+
+# Problem templates already handle basic bounds via the trajectory:
+
+prob = UnitarySmoothPulseProblem(sys, U_goal, N)
+
+# Control bounds come from `sys.drive_bounds`:
+println("Control bounds from system: ", sys.drive_bounds)
+
+# Derivative bounds specified in problem:
+prob_with_bounds = UnitarySmoothPulseProblem(sys, U_goal, N;
+    du_bound=0.5,    # First derivative bound
+    ddu_bound=0.1    # Second derivative bound
+)
+
+# ### Custom Variable Bounds
+
+# To add bounds on other trajectory variables:
+
+# Create trajectory with custom bounds
+traj = initialize_unitary_trajectory(
+    U_goal,
+    N,
+    sys.T_max / N,
+    sys.n_drives,
+    (sys.drive_bounds, fill((-0.5, 0.5), sys.n_drives), fill((-0.1, 0.1), sys.n_drives));
+    state_name=:Ũ⃗,
+    control_name=:u,
+)
+
+# Add a custom bounded variable
+traj_extended = NamedTrajectory(
+    merge(traj.data, Dict(:custom_var => randn(1, N))),
+    controls=traj.controls,
+    timestep=traj.timestep,
+    bounds=merge(traj.bounds, Dict(:custom_var => [(-1.0, 1.0)])),
+)
+
+# ## 2. Linear Equality Constraints
+
+# Use `EqualityConstraint` for linear constraints of form `A*z = b`:
+
+using LinearAlgebra
+
+# Example: Force controls to sum to zero at each timestep
+# (useful for differential drive schemes)
+
+prob = UnitarySmoothPulseProblem(sys, U_goal, N)
+
+# Get trajectory dimensions
+n_drives = sys.n_drives
+n_timesteps = N
+
+# Create constraint: u1(t) + u2(t) = 0 for all t
+constraints = AbstractConstraint[]
+
+for k in 1:n_timesteps
+    # Get indices for u1 and u2 at timestep k
+    u1_idx = (k-1) * n_drives + 1
+    u2_idx = (k-1) * n_drives + 2
+    
+    # A * [u1; u2] = 0, where A = [1, 1]
+    A = zeros(1, prob.trajectory.dim * N)
+    A[1, u1_idx] = 1.0
+    A[1, u2_idx] = 1.0
+    
+    constraint = EqualityConstraint(A, [0.0], [k])
+    push!(constraints, constraint)
+end
+
+# Create problem with constraints
+prob_with_constraints = DirectTrajOptProblem(
+    prob.trajectory,
+    prob.objective,
+    prob.dynamics;
+    constraints=vcat(prob.constraints, constraints)
+)
+
+# ## 3. Nonlinear Path Constraints
+
+# For nonlinear constraints applied at each timestep, use `NonlinearKnotPointConstraint`:
+
+# Example: Limit total control power: √(u₁² + u₂²) ≤ P_max
+
+P_max = 0.8
+
+# Define constraint function
+function power_constraint(Z⃗, k, traj)
+    # Extract controls at timestep k
+    u_slice = slice(k, :u, traj.dims)
+    u = Z⃗[u_slice]
+    
+    # Return constraint value (should be ≤ 0)
+    return sqrt(sum(u.^2)) - P_max
+end
+
+# Create the constraint
+power_constr = NonlinearKnotPointConstraint(
+    power_constraint,
+    prob.trajectory,
+    1:N;  # Apply at all timesteps
+    name=:power_limit
+)
+
+# Add to problem
+prob_power = DirectTrajOptProblem(
+    prob.trajectory,
+    prob.objective,
+    prob.dynamics;
+    constraints=vcat(prob.constraints, [power_constr])
+)
+
+# ## 4. Nonlinear Global Constraints
+
+# For constraints over the entire trajectory, use `NonlinearGlobalConstraint`:
+
+# Example: Limit total energy: ∫ u²(t) dt ≤ E_budget
+
+E_budget = 10.0
+
+function energy_constraint(Z⃗, traj)
+    total_energy = 0.0
+    
+    for k in 1:traj.T
+        u_slice = slice(k, :u, traj.dims)
+        Δt_slice = slice(k, :Δt, traj.dims)
+        
+        u = Z⃗[u_slice]
+        Δt = Z⃗[Δt_slice][1]
+        
+        total_energy += sum(u.^2) * Δt
+    end
+    
+    return total_energy - E_budget
+end
+
+energy_constr = NonlinearGlobalConstraint(
+    energy_constraint,
+    prob.trajectory;
+    name=:energy_budget
+)
+
+# ## 5. Leakage Suppression (Multilevel Systems)
+
+# For multilevel systems, suppress population in leakage levels using `PiccoloOptions`:
+
+# Define a 3-level system (qubit + 1 leakage level)
+a = annihilate(3)
+H_drives_3level = [(a + a')/2, (a - a')/(2im)]
+sys_3level = QuantumSystem(H_drives_3level, 10.0, [1.0, 1.0])
+
+# Target: X gate in computational subspace
+U_goal_embedded = EmbeddedOperator(GATES.X, sys_3level)
+
+# Enable leakage suppression
+opts = PiccoloOptions(
+    leakage_constraint=true,
+    leakage_constraint_value=1e-3,  # Maximum leakage population
+    leakage_cost=1.0                 # Penalty weight
+)
+
+prob_leakage = UnitarySmoothPulseProblem(
+    sys_3level,
+    U_goal_embedded,
+    N;
+    piccolo_options=opts
+)
+
+solve!(prob_leakage; max_iter=100)
+
+# Check leakage:
+# (Leakage evaluation code would go here)
+
+# ## 6. Soft Constraints vs Hard Constraints
+
+# **Hard constraints** (equality/inequality):
+# - Must be satisfied exactly
+# - Can make problem infeasible
+# - Use for physics requirements
+
+# **Soft constraints** (penalties in objective):
+# - Violated at a cost
+# - Always feasible
+# - Use for preferences
+
+# Example: Add soft constraint via objective penalty
+
+# Soft control smoothness (already built-in via R_u, R_du, R_ddu)
+prob_smooth = UnitarySmoothPulseProblem(sys, U_goal, N;
+    R_u=0.01,    # Control amplitude penalty
+    R_du=0.1,    # First derivative penalty (smoothness)
+    R_ddu=1.0    # Second derivative penalty (extra smooth)
+)
+
+# ## 7. State Constraints (Avoid Forbidden Regions)
+
+# Example: Keep qubit in upper hemisphere of Bloch sphere
+
+function upper_hemisphere_constraint(Z⃗, k, traj)
+    # Extract state at timestep k
+    state_slice = slice(k, :Ũ⃗, traj.dims)
+    Ũ⃗ = Z⃗[state_slice]
+    
+    # Convert to unitary (implementation depends on representation)
+    # For isomorphic vectors, extract Z expectation value
+    # Here's conceptual code - actual implementation depends on system
+    
+    # Return constraint: should keep ⟨Z⟩ > 0
+    # return -Z_expectation  # Example
+    
+    return 0.0  # Placeholder
+end
+
+# ## 8. Time-Dependent Constraints
+
+# Example: Ramp control amplitude over time
+
+function amplitude_ramp_constraint(Z⃗, k, traj)
+    # Maximum amplitude increases linearly with time
+    max_amp = 0.5 * (k / traj.T)  # Ramp from 0 to 0.5
+    
+    u_slice = slice(k, :u, traj.dims)
+    u = Z⃗[u_slice]
+    
+    # All controls should be below max_amp
+    return maximum(abs.(u)) - max_amp
+end
+
+ramp_constr = NonlinearKnotPointConstraint(
+    amplitude_ramp_constraint,
+    prob.trajectory,
+    1:N;
+    name=:amplitude_ramp
+)
+
+# ## 9. Coupled State-Control Constraints
+
+# Example: Control amplitude proportional to state fidelity
+# (more aggressive control when far from target)
+
+function adaptive_control_constraint(Z⃗, k, traj)
+    state_slice = slice(k, traj.state, traj.dims)
+    u_slice = slice(k, :u, traj.dims)
+    
+    # Get distance from target (simplified)
+    # state = Z⃗[state_slice]
+    # distance = ... (compute distance metric)
+    
+    u = Z⃗[u_slice]
+    
+    # Constraint: |u| ≤ f(distance)
+    # return maximum(abs.(u)) - constraint_function(distance)
+    
+    return 0.0  # Placeholder
+end
+
+# ## Key Takeaways
+
+# 1. **Start with problem templates** - they include sensible defaults
+# 2. **Use PiccoloOptions** for common constraints (leakage, robustness)
+# 3. **EqualityConstraint** for linear constraints
+# 4. **NonlinearKnotPointConstraint** for per-timestep nonlinear constraints
+# 5. **NonlinearGlobalConstraint** for trajectory-wide constraints
+# 6. **Soft constraints** (objectives) are more robust than hard constraints
+# 7. **Test feasibility** - start with loose constraints, then tighten
+#
+# ## Debugging Constraints
+#
+# If your problem becomes infeasible:
+# 1. Check constraint values before optimization
+# 2. Relax bounds temporarily
+# 3. Use soft constraints instead of hard
+# 4. Verify constraint gradients (use finite differences)
+# 5. Start from a feasible initial trajectory
+#
+# ## Next Steps
+#
+# - [Custom Objectives](@ref) - Add custom cost functions
+# - [Initial Trajectories](@ref) - Better initialization strategies
+# - [PiccoloOptions Reference](@ref) - Built-in constraint options
diff --git a/docs/literate/man/custom_objectives.jl b/docs/literate/man/custom_objectives.jl
new file mode 100644
index 00000000..d60f65c6
--- /dev/null
+++ b/docs/literate/man/custom_objectives.jl
@@ -0,0 +1,345 @@
+# # Custom Objectives and Regularization
+
+# Beyond the default fidelity objective, you often want to customize the cost function to achieve specific control properties. This guide shows how to tune regularization weights and add custom objectives.
+
+using QuantumCollocation
+using PiccoloQuantumObjects
+using NamedTrajectories
+using DirectTrajOpt
+
+# ## Basic Setup
+
+H_drift = 0.1 * PAULIS.Z
+H_drives = [PAULIS.X, PAULIS.Y]
+sys = QuantumSystem(H_drift, H_drives, 10.0, [1.0, 1.0])
+U_goal = GATES.H
+N = 51
+
+# ## Understanding Default Objectives
+
+# A standard `UnitarySmoothPulseProblem` minimizes:
+# ```math
+# J = Q \cdot \text{infidelity} + \frac{1}{2}\sum_k (R_u u_k^2 + R_{du} \dot{u}_k^2 + R_{ddu} \ddot{u}_k^2)
+# ```
+
+# Let's see the defaults:
+
+prob_default = UnitarySmoothPulseProblem(sys, U_goal, N)
+
+# Default weights (from problem template):
+# - Q = 100.0 (infidelity weight)
+# - R_u = 0.01 (control amplitude)
+# - R_du = 0.01 (first derivative)  
+# - R_ddu = 0.01 (second derivative)
+
+# ## 1. Tuning Regularization Weights
+
+# ### Infidelity Weight (Q)
+
+# **Higher Q** → prioritize fidelity over smoothness
+# **Lower Q** → prioritize smooth pulses over fidelity
+
+prob_high_fidelity = UnitarySmoothPulseProblem(sys, U_goal, N;
+    Q=1000.0,    # 10× default - demand high fidelity
+    R_u=0.01,
+    R_du=0.01,
+    R_ddu=0.01
+)
+
+solve!(prob_high_fidelity; max_iter=100)
+println("High Q fidelity: ", unitary_rollout_fidelity(prob_high_fidelity.trajectory, sys))
+
+prob_smooth_priority = UnitarySmoothPulseProblem(sys, U_goal, N;
+    Q=10.0,      # Lower Q - accept lower fidelity for smoother pulses
+    R_u=0.1,
+    R_du=1.0,
+    R_ddu=10.0
+)
+
+solve!(prob_smooth_priority; max_iter=100)
+println("Smooth priority fidelity: ", unitary_rollout_fidelity(prob_smooth_priority.trajectory, sys))
+
+# ### Control Amplitude (R_u)
+
+# Penalizes large control values: minimizes average power
+
+# Low amplitude preference:
+prob_low_power = UnitarySmoothPulseProblem(sys, U_goal, N;
+    R_u=1.0      # 100× default - strong amplitude penalty
+)
+
+solve!(prob_low_power; max_iter=100)
+
+# Check power:
+avg_power_default = mean(prob_default.trajectory.u.^2)
+avg_power_low = mean(prob_low_power.trajectory.u.^2)
+println("Average power - default: ", avg_power_default)
+println("Average power - low: ", avg_power_low)
+
+# ### First Derivative (R_du)
+
+# Penalizes rapid changes: enforces smoothness
+
+# Very smooth pulses:
+prob_very_smooth = UnitarySmoothPulseProblem(sys, U_goal, N;
+    R_du=10.0    # Strong smoothness penalty
+)
+
+solve!(prob_very_smooth; max_iter=100)
+
+# ### Second Derivative (R_ddu)
+
+# Penalizes acceleration: prevents sharp corners
+
+# Ultra-smooth (C² continuous) pulses:
+prob_ultra_smooth = UnitarySmoothPulseProblem(sys, U_goal, N;
+    R_ddu=100.0  # Penalize curvature heavily
+)
+
+solve!(prob_ultra_smooth; max_iter=100)
+
+# ## 2. Per-Drive Regularization
+
+# Different drives may need different regularization:
+
+# Example: Strong Y drive, gentle X drive
+R_u_per_drive = [0.01, 0.1]    # Penalize Y more
+R_du_per_drive = [0.01, 1.0]   # Y must be smoother
+R_ddu_per_drive = [0.01, 10.0]
+
+prob_asymmetric = UnitarySmoothPulseProblem(sys, U_goal, N;
+    R_u=R_u_per_drive,
+    R_du=R_du_per_drive,
+    R_ddu=R_ddu_per_drive
+)
+
+solve!(prob_asymmetric; max_iter=100)
+
+# ## 3. Adding Custom Objectives
+
+# To add completely custom cost functions, work with the `Objective` type.
+
+# ### Example: Minimize Peak Power
+
+# Add a penalty for maximum control amplitude:
+
+using ForwardDiff
+
+# Define a custom objective function
+function peak_power_objective(Z⃗, traj)
+    cost = 0.0
+    for k in 1:traj.T
+        u_slice = slice(k, :u, traj.dims)
+        u = Z⃗[u_slice]
+        # Smooth approximation of max: log-sum-exp
+        cost += log(sum(exp.(10 .* u.^2))) / 10
+    end
+    return cost
+end
+
+# Compute gradient (for optimization)
+function peak_power_gradient(Z⃗, traj)
+    return ForwardDiff.gradient(z -> peak_power_objective(z, traj), Z⃗)
+end
+
+# Compute Hessian structure (sparsity pattern)
+function peak_power_hessian_structure(traj)
+    # For simple objectives, dense structure is fine
+    n = traj.dim * traj.T
+    return [(i, j) for i in 1:n for j in 1:i]
+end
+
+# Compute Hessian values
+function peak_power_hessian(Z⃗, traj)
+    H = ForwardDiff.hessian(z -> peak_power_objective(z, traj), Z⃗)
+    # Extract lower triangle
+    vals = Float64[]
+    for i in 1:size(H, 1)
+        for j in 1:i
+            push!(vals, H[i, j])
+        end
+    end
+    return vals
+end
+
+# Create Objective
+peak_power_obj = Objective(
+    peak_power_objective,
+    peak_power_gradient,
+    peak_power_hessian,
+    peak_power_hessian_structure
+)
+
+# Combine with standard objective
+prob_with_peak = UnitarySmoothPulseProblem(sys, U_goal, N)
+total_objective = prob_with_peak.objective + 0.1 * peak_power_obj
+
+# Create new problem with combined objective
+prob_custom = DirectTrajOptProblem(
+    prob_with_peak.trajectory,
+    total_objective,
+    prob_with_peak.dynamics;
+    constraints=prob_with_peak.constraints
+)
+
+solve!(prob_custom; max_iter=100)
+
+# ## 4. Time-Dependent Weights
+
+# Example: Penalize control more at the end (for clean turn-off)
+
+function time_weighted_control_objective(Z⃗, traj)
+    cost = 0.0
+    for k in 1:traj.T
+        # Weight increases linearly with time
+        weight = k / traj.T
+        
+        u_slice = slice(k, :u, traj.dims)
+        u = Z⃗[u_slice]
+        
+        cost += weight * sum(u.^2)
+    end
+    return cost
+end
+
+# (Similar gradient/hessian computation as above)
+
+# ## 5. Energy Budget Objective
+
+# Instead of just penalizing amplitude, minimize total energy:
+
+function energy_objective(Z⃗, traj)
+    energy = 0.0
+    for k in 1:traj.T
+        u_slice = slice(k, :u, traj.dims)
+        Δt_slice = slice(k, :Δt, traj.dims)
+        
+        u = Z⃗[u_slice]
+        Δt = Z⃗[Δt_slice][1]
+        
+        energy += sum(u.^2) * Δt
+    end
+    return energy
+end
+
+# ## 6. Spectral Regularization
+
+# Penalize high-frequency components (alternative to derivative penalties):
+
+using FFTW
+
+function spectral_regularization(Z⃗, traj)
+    cost = 0.0
+    for drive_idx in 1:traj.dims[:u]
+        # Extract time series for this drive
+        u_timeseries = [Z⃗[slice(k, :u, traj.dims)[drive_idx]] for k in 1:traj.T]
+        
+        # FFT
+        spectrum = fft(u_timeseries)
+        
+        # Penalize high frequencies (e.g., > 50% Nyquist)
+        cutoff = div(length(spectrum), 2)
+        high_freq = spectrum[cutoff:end]
+        
+        cost += sum(abs2.(high_freq))
+    end
+    return cost
+end
+
+# ## 7. Robustness Objectives
+
+# Penalize sensitivity to parameters:
+
+# Example: Minimize gradient w.r.t. system parameters
+function robustness_objective(Z⃗, traj, sys)
+    # Compute fidelity gradient w.r.t. drift Hamiltonian
+    # (Requires variational quantum system)
+    # This is conceptual - see UnitaryVariationalProblem for implementation
+    return 0.0  # Placeholder
+end
+
+# ## 8. Physical Realism Objectives
+
+# Example: Match target pulse shape (e.g., Gaussian envelope)
+
+function gaussian_envelope_objective(Z⃗, traj)
+    cost = 0.0
+    for k in 1:traj.T
+        # Desired Gaussian envelope
+        t_norm = (k - traj.T/2) / (traj.T/4)
+        target_amplitude = exp(-t_norm^2)
+        
+        u_slice = slice(k, :u, traj.dims)
+        u = Z⃗[u_slice]
+        
+        # Penalize deviation from Gaussian envelope
+        cost += (norm(u) - target_amplitude)^2
+    end
+    return cost
+end
+
+# ## 9. Combining Multiple Objectives
+
+# Objectives are additive - you can combine many:
+
+prob = UnitarySmoothPulseProblem(sys, U_goal, N;
+    Q=100.0,      # High fidelity
+    R_u=0.01,     # Moderate amplitude penalty
+    R_du=0.1,     # Strong smoothness
+    R_ddu=1.0     # Very strong curvature penalty
+)
+
+# Then add custom objectives:
+# total_obj = prob.objective + 0.1 * energy_obj + 0.01 * spectral_obj + ...
+
+# ## Common Regularization Recipes
+
+# **Smooth, implementable pulses** (most common):
+# ```julia
+# R_u=0.01, R_du=0.1, R_ddu=1.0
+# ```
+
+# **Low power, accept longer time**:
+# ```julia
+# R_u=1.0, R_du=0.01, R_ddu=0.01
+# ```
+
+# **Ultra-smooth for noisy hardware**:
+# ```julia
+# R_u=0.01, R_du=1.0, R_ddu=10.0
+# ```
+
+# **High fidelity at any cost**:
+# ```julia
+# Q=1000.0, R_u=1e-6, R_du=1e-6, R_ddu=1e-6
+# ```
+
+# **Aggressive pulses for fast gates**:
+# ```julia
+# Q=100.0, R_u=1e-4, R_du=1e-4, R_ddu=1e-3
+# ```
+
+# ## Tuning Strategy
+
+# 1. **Start with defaults** - they work well for most cases
+# 2. **Adjust one weight at a time** - see the effect
+# 3. **Check convergence** - some combinations are hard to optimize
+# 4. **Validate physically** - can your hardware actually do this?
+# 5. **Iterate** - optimal weights are problem-specific
+
+# ## Key Takeaways
+
+# - `Q` controls fidelity vs smoothness tradeoff
+# - `R_u` penalizes control amplitude (power)
+# - `R_du` enforces smoothness (rate limits)
+# - `R_ddu` prevents sharp corners (acceleration limits)
+# - Per-drive weights allow asymmetric regularization
+# - Custom objectives enable arbitrary cost functions
+# - Start with defaults, tune empirically
+# - Higher regularization → smoother but slower convergence
+#
+# ## Next Steps
+#
+# - [Adding Constraints](@ref) - Enforce hard limits
+# - [Initial Trajectories](@ref) - Better starting points
+# - [Problem Templates Overview](@ref) - Choose the right template
diff --git a/docs/literate/man/initial_trajectories.jl b/docs/literate/man/initial_trajectories.jl
new file mode 100644
index 00000000..0aaf0498
--- /dev/null
+++ b/docs/literate/man/initial_trajectories.jl
@@ -0,0 +1,341 @@
+# # Initial Trajectory Strategies
+
+# Good initialization is crucial for nonlinear optimization. A well-chosen initial trajectory can:
+# - Reduce solve time by 10-100×
+# - Avoid local minima
+# - Improve convergence reliability
+# - Enable harder problems to solve
+
+# This guide shows various initialization strategies for quantum optimal control problems.
+
+using QuantumCollocation
+using PiccoloQuantumObjects
+using NamedTrajectories
+
+# ## Basic Setup
+
+H_drift = 0.1 * PAULIS.Z
+H_drives = [PAULIS.X, PAULIS.Y]
+sys = QuantumSystem(H_drift, H_drives, 10.0, [1.0, 1.0])
+U_goal = GATES.H
+N = 51
+
+# ## 1. Default Initialization (Random + Rollout)
+
+# Problem templates use default initialization:
+
+prob_default = UnitarySmoothPulseProblem(sys, U_goal, N)
+
+# Check initial fidelity:
+fid_init_default = unitary_rollout_fidelity(prob_default.trajectory, sys)
+println("Default initialization fidelity: ", fid_init_default)
+
+# Default process:
+# 1. Random controls (small amplitude)
+# 2. Forward rollout to get states
+# 3. Satisfy derivative relationships
+
+# ## 2. Warm Start from Previous Solution
+
+# **Best strategy:** Use solution from similar problem
+
+# Solve a simpler problem first:
+prob_easy = UnitarySmoothPulseProblem(sys, U_goal, N;
+    Q=10.0,      # Lower fidelity requirement
+    R_u=0.001    # Less regularization
+)
+solve!(prob_easy; max_iter=50)
+
+# Use as initialization for harder problem:
+prob_hard = UnitarySmoothPulseProblem(sys, U_goal, N;
+    Q=1000.0,    # Higher fidelity
+    R_u=0.1,     # More regularization
+    init_trajectory=prob_easy.trajectory
+)
+solve!(prob_hard; max_iter=100)
+
+println("Warm start fidelity: ", unitary_rollout_fidelity(prob_hard.trajectory, sys))
+
+# ## 3. Custom Control Guess (u_guess)
+
+# Provide initial control pulses based on physical intuition:
+
+# Example: Smooth ramp-up/ramp-down pulse
+
+function create_smooth_pulse(N, n_drives, max_amplitude=0.5)
+    u = zeros(n_drives, N)
+    
+    for k in 1:N
+        # Smooth envelope: sin² ramp
+        t_norm = (k - 1) / (N - 1)
+        envelope = sin(π * t_norm)^2
+        
+        # Drive 1: cosine modulation
+        u[1, k] = max_amplitude * envelope * cos(2π * t_norm)
+        
+        # Drive 2: sine modulation (π/2 out of phase)
+        u[2, k] = max_amplitude * envelope * sin(2π * t_norm)
+    end
+    
+    return u
+end
+
+u_guess = create_smooth_pulse(N, sys.n_drives)
+
+prob_with_guess = UnitarySmoothPulseProblem(sys, U_goal, N;
+    u_guess=u_guess
+)
+
+fid_init_guess = unitary_rollout_fidelity(prob_with_guess.trajectory, sys)
+println("Custom guess fidelity: ", fid_init_guess)
+
+solve!(prob_with_guess; max_iter=100)
+
+# ## 4. Geodesic Initialization
+
+# For unitary gates, initialize state trajectory along geodesic path:
+
+prob_geodesic = UnitarySmoothPulseProblem(sys, U_goal, N;
+    piccolo_options=PiccoloOptions(geodesic=true)
+)
+
+fid_init_geodesic = unitary_rollout_fidelity(prob_geodesic.trajectory, sys)
+println("Geodesic initialization fidelity: ", fid_init_geodesic)
+
+solve!(prob_geodesic; max_iter=100)
+
+# Geodesic initialization:
+# - Interpolates between identity and target unitary
+# - Follows shortest path in SU(N)
+# - Often gives better starting fidelity
+# - Particularly useful for large gates
+
+# ## 5. GRAPE-style Initialization
+
+# Use constant pulses at maximum amplitude (aggressive):
+
+u_grape = ones(sys.n_drives, N) .* 0.8  # 80% of maximum
+u_grape[1, :] .*= cos.(range(0, 2π, length=N))  # Modulate drive 1
+u_grape[2, :] .*= sin.(range(0, 2π, length=N))  # Modulate drive 2
+
+prob_grape = UnitarySmoothPulseProblem(sys, U_goal, N;
+    u_guess=u_grape
+)
+
+solve!(prob_grape; max_iter=100)
+
+# GRAPE-style works well when:
+# - Strong controls available
+# - Short time horizons
+# - Seeking maximum speed
+
+# ## 6. Time Scaling / Trajectory Resampling
+
+# Solve at coarse resolution, then refine:
+
+# Coarse problem (fewer timesteps)
+N_coarse = 25
+prob_coarse = UnitarySmoothPulseProblem(sys, U_goal, N_coarse)
+solve!(prob_coarse; max_iter=100)
+
+# Resample to finer resolution
+function resample_trajectory(traj_coarse, N_fine)
+    # Linear interpolation of controls
+    N_coarse = traj_coarse.T
+    u_coarse = traj_coarse.u
+    
+    u_fine = zeros(size(u_coarse, 1), N_fine)
+    for drive_idx in 1:size(u_coarse, 1)
+        # Simple linear interpolation
+        t_coarse = range(0, 1, length=N_coarse)
+        t_fine = range(0, 1, length=N_fine)
+        
+        for (i, t) in enumerate(t_fine)
+            # Find surrounding coarse points
+            k = min(N_coarse - 1, max(1, floor(Int, t * (N_coarse - 1)) + 1))
+            α = (t * (N_coarse - 1)) - (k - 1)
+            
+            u_fine[drive_idx, i] = (1 - α) * u_coarse[drive_idx, k] + 
+                                    α * u_coarse[drive_idx, min(k + 1, N_coarse)]
+        end
+    end
+    
+    return u_fine
+end
+
+u_fine = resample_trajectory(prob_coarse.trajectory, N)
+
+# Solve fine problem with resampled initialization
+prob_fine = UnitarySmoothPulseProblem(sys, U_goal, N;
+    u_guess=u_fine
+)
+solve!(prob_fine; max_iter=100)
+
+println("Coarse-to-fine final fidelity: ", unitary_rollout_fidelity(prob_fine.trajectory, sys))
+
+# ## 7. Adiabatic / Homotopy Continuation
+
+# Gradually increase difficulty:
+
+# Strategy: Solve sequence with increasing Q (fidelity weight)
+Q_sequence = [1.0, 10.0, 100.0, 1000.0]
+prob_current = nothing
+
+for Q in Q_sequence
+    if isnothing(prob_current)
+        prob_current = UnitarySmoothPulseProblem(sys, U_goal, N; Q=Q)
+    else
+        prob_current = UnitarySmoothPulseProblem(sys, U_goal, N;
+            Q=Q,
+            init_trajectory=prob_current.trajectory
+        )
+    end
+    
+    solve!(prob_current; max_iter=50)
+    fid = unitary_rollout_fidelity(prob_current.trajectory, sys)
+    println("Q=$Q → fidelity: $fid")
+end
+
+# ## 8. Random Restart Strategy
+
+# Try multiple random initializations, keep best:
+
+function solve_with_random_restarts(sys, U_goal, N, n_restarts=5)
+    best_fidelity = 0.0
+    best_trajectory = nothing
+    
+    for restart in 1:n_restarts
+        prob = UnitarySmoothPulseProblem(sys, U_goal, N)
+        solve!(prob; max_iter=50, verbose=false)
+        
+        fid = unitary_rollout_fidelity(prob.trajectory, sys)
+        println("Restart $restart: fidelity = $fid")
+        
+        if fid > best_fidelity
+            best_fidelity = fid
+            best_trajectory = prob.trajectory
+        end
+    end
+    
+    return best_trajectory
+end
+
+# Run random restarts
+best_traj = solve_with_random_restarts(sys, U_goal, N, 3)
+
+# Refine best solution
+prob_refined = UnitarySmoothPulseProblem(sys, U_goal, N;
+    init_trajectory=best_traj
+)
+solve!(prob_refined; max_iter=100)
+
+# ## 9. Physical Insights: Bang-Bang to Smooth
+
+# Start with bang-bang (on/off) controls, smooth them out:
+
+function create_bangbang_pulse(N, n_drives, switch_times)
+    u = zeros(n_drives, N)
+    
+    for (drive_idx, switches) in enumerate(switch_times)
+        state = 1.0
+        switch_idx = 1
+        
+        for k in 1:N
+            if switch_idx <= length(switches) && k >= switches[switch_idx]
+                state *= -1
+                switch_idx += 1
+            end
+            u[drive_idx, k] = state * 0.8
+        end
+    end
+    
+    return u
+end
+
+# Define switching pattern
+switch_times = [[10, 25, 40], [5, 30, 45]]
+u_bangbang = create_bangbang_pulse(N, sys.n_drives, switch_times)
+
+prob_bangbang = UnitarySmoothPulseProblem(sys, U_goal, N;
+    u_guess=u_bangbang,
+    R_du=10.0,    # Strong derivative penalty will smooth it
+    R_ddu=100.0
+)
+
+solve!(prob_bangbang; max_iter=200)
+
+# ## 10. Transfer Learning from Similar Gates
+
+# Use solution from different but related gate:
+
+# Solve for X gate
+prob_X = UnitarySmoothPulseProblem(sys, GATES.X, N)
+solve!(prob_X; max_iter=100)
+
+# Use as initialization for Y gate (related by π/2 rotation)
+prob_Y = UnitarySmoothPulseProblem(sys, GATES.Y, N;
+    init_trajectory=prob_X.trajectory
+)
+solve!(prob_Y; max_iter=100)
+
+println("Transfer learning Y-gate fidelity: ", unitary_rollout_fidelity(prob_Y.trajectory, sys))
+
+# ## Initialization Strategy Decision Tree
+
+# **Easy problem** (single qubit, strong controls):
+# → Default initialization works fine
+
+# **Medium problem** (two qubits, moderate controls):
+# → Geodesic initialization OR coarse-to-fine
+
+# **Hard problem** (many levels, weak controls, tight constraints):
+# → Homotopy continuation (gradually increase difficulty)
+# → OR warm start from related solved problem
+# → OR coarse-to-fine with multiple refinements
+
+# **Unknown difficulty**:
+# → Try default first
+# → If convergence slow: switch to geodesic
+# → If still slow: use coarse-to-fine
+# → If repeatedly failing: random restarts + homotopy
+
+# ## Best Practices
+
+# 1. **Always check initial fidelity** - should be > 0.01 for most problems
+# 2. **Warm start when possible** - reuse previous solutions
+# 3. **Scale up gradually** - coarse to fine, easy to hard
+# 4. **Save successful trajectories** - build a library
+# 5. **Use physical intuition** - u_guess can encode domain knowledge
+# 6. **Be patient with hard problems** - good initialization takes time
+# 7. **Monitor convergence** - poor init shows in first few iterations
+
+# ## Common Issues and Fixes
+
+# **Issue**: Optimizer gives up immediately (max_iter reached, low fidelity)
+# **Fix**: Better initialization - try geodesic or warm start
+
+# **Issue**: Solution oscillates, doesn't converge
+# **Fix**: Increase regularization (R_du, R_ddu) OR use smoother u_guess
+
+# **Issue**: Infeasible problem from start
+# **Fix**: Relax constraints temporarily, find feasible point, then tighten
+
+# **Issue**: Good fidelity but violates constraints
+# **Fix**: Start from feasible trajectory respecting all constraints
+
+# ## Key Takeaways
+
+# - Default initialization works for ~80% of problems
+# - Warm starting from related solutions is the best strategy
+# - Geodesic initialization helps for large unitary gates
+# - Coarse-to-fine scaling is robust for hard problems
+# - Physical intuition (u_guess) can dramatically help
+# - Random restarts catch lucky initializations
+# - Homotopy continuation solves hard problems reliably
+# - Save and reuse successful trajectories
+#
+# ## Next Steps
+#
+# - [Custom Objectives](@ref) - Tune cost function
+# - [Adding Constraints](@ref) - Enforce requirements
+# - [Problem Templates Overview](@ref) - Choose right template
diff --git a/docs/literate/man/ket_problem_templates.jl b/docs/literate/man/ket_problem_templates.jl
index e25d22f3..c51fa21b 100644
--- a/docs/literate/man/ket_problem_templates.jl
+++ b/docs/literate/man/ket_problem_templates.jl
@@ -20,28 +20,25 @@ intial state, `ψ_init`, to a target state, `ψ_goal`.
 =#
 
 # _define the quantum system_
-T_max = 1.0
-u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-system = QuantumSystem(0.1 * PAULIS.Z, [PAULIS.X, PAULIS.Y], T_max, u_bounds)
+system = QuantumSystem(0.1 * PAULIS.Z, [PAULIS.X, PAULIS.Y], 10.0, [1.0, 1.0])
 ψ_init = Vector{ComplexF64}([1.0, 0.0])
 ψ_goal = Vector{ComplexF64}([0.0, 1.0])
 N = 51
-Δt = 0.2
 
 # _create the smooth pulse problem_
-state_prob = QuantumStateSmoothPulseProblem(system, ψ_init, ψ_goal, N, Δt);
+state_prob = QuantumStateSmoothPulseProblem(system, ψ_init, ψ_goal, N);
 
 # _check the fidelity before solving_
 println("Before: ", rollout_fidelity(state_prob.trajectory, system))
 
 # _solve the problem_
-load_path = joinpath(dirname(Base.active_project()), "data/ket_problem_templates_state_prob_89ee72.jld2") # hide
-state_prob.trajectory = load_traj(load_path) # hide
-nothing # hide
+# load_path = joinpath(dirname(Base.active_project()), "data/ket_problem_templates_state_prob_89ee72.jld2") # hide
+# state_prob.trajectory = load_traj(load_path) # hide
+# nothing # hide
+# solve!(state_prob, max_iter=100);
 
 #=
 ```julia
-solve!(state_prob, max_iter=100, verbose=true, print_level=1);
 ```
 
 ```@raw html
@@ -81,13 +78,13 @@ min_state_prob = QuantumStateMinimumTimeProblem(state_prob, ψ_goal);
 println("Duration before: ", get_duration(state_prob.trajectory))
 
 # _solve the minimum time problem_
-load_path = joinpath(dirname(Base.active_project()), "data/ket_problem_templates_state_prob_min_time_89ee72.jld2") # hide
-min_state_prob.trajectory = load_traj(load_path) # hide
-nothing # hide
+# load_path = joinpath(dirname(Base.active_project()), "data/ket_problem_templates_state_prob_min_time_89ee72.jld2") # hide
+# min_state_prob.trajectory = load_traj(load_path) # hide
+# nothing # hide
+# solve!(min_state_prob, max_iter=100);
 
 #=
 ```julia
-solve!(min_state_prob, max_iter=100, verbose=true, print_level=1);
 ```
 
 ```@raw html
@@ -122,22 +119,20 @@ QuantumStateSamplingProblem
 =#
 
 # _create a sampling problem_
-T_max = 1.0
-u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-driftless_system = QuantumSystem([PAULIS.X, PAULIS.Y], T_max, u_bounds)
-sampling_state_prob = QuantumStateSamplingProblem([system, driftless_system], ψ_init, ψ_goal, N, Δt);
+driftless_system = QuantumSystem([PAULIS.X, PAULIS.Y], 10.0, [1.0, 1.0])
+sampling_state_prob = QuantumStateSamplingProblem([system, driftless_system], ψ_init, ψ_goal, N);
 
 # _new keys are added to the trajectory for the new states_
 println(sampling_state_prob.trajectory.state_names)
 
 # _solve the sampling problem for a few iterations_
-load_path = joinpath(dirname(Base.active_project()), "data/ket_problem_templates_sampling_state_89ee72.jld2") # hide
-sampling_state_prob.trajectory = load_traj(load_path) # hide
-nothing # hide
+# load_path = joinpath(dirname(Base.active_project()), "data/ket_problem_templates_sampling_state_89ee72.jld2") # hide
+# sampling_state_prob.trajectory = load_traj(load_path) # hide
+# nothing # hide
+# solve!(sampling_state_prob, max_iter=25);
 
 #=
 ```julia
-solve!(sampling_state_prob, max_iter=25, verbose=true, print_level=1);
 ```
 
 ```@raw html
diff --git a/docs/literate/man/piccolo_options.jl b/docs/literate/man/piccolo_options.jl
new file mode 100644
index 00000000..f477b8de
--- /dev/null
+++ b/docs/literate/man/piccolo_options.jl
@@ -0,0 +1,197 @@
+# # PiccoloOptions Reference
+
+# `PiccoloOptions` provides advanced configuration for problem templates. This page documents all available options and their effects.
+
+using QuantumCollocation
+using PiccoloQuantumObjects
+
+# ## Creating PiccoloOptions
+
+# Default options:
+opts = PiccoloOptions()
+
+# Custom options:
+opts_custom = PiccoloOptions(
+    verbose = true,
+    leakage_constraint = true,
+    bound_state = true
+)
+
+# Pass to any problem template:
+system = QuantumSystem(0.1 * PAULIS.Z, [PAULIS.X, PAULIS.Y], 10.0, [1.0, 1.0])
+U_goal = EmbeddedOperator(GATES.H, system)
+N = 51
+
+prob = UnitarySmoothPulseProblem(
+    system, U_goal, N;
+    piccolo_options = opts_custom
+)
+
+# ## General Options
+
+# ### `verbose::Bool = false`
+# Print detailed information during problem setup.
+
+opts_verbose = PiccoloOptions(verbose = true)
+# Shows: constraint counts, objective terms, trajectory initialization details
+
+# ### `free_time::Bool = false`
+# Allow timesteps to vary (used internally by minimum time problems).
+# Typically not set directly by users.
+
+# ## State and Unitary Options
+
+# ### `bound_state::Bool = false`
+# Constrain state/unitary to lie on unit sphere (for numerical stability).
+# **Recommended for difficult convergence.**
+
+opts_bounded = PiccoloOptions(bound_state = true)
+# Adds constraint: ||ψ||² = 1 (ket) or ||U||² = N (unitary)
+
+# ### `geodesic::Bool = true`
+# Use geodesic interpolation for initial trajectory.
+# - `true`: States evolve along shortest path on manifold (better initial guess)
+# - `false`: Linear interpolation (simpler but often worse)
+
+opts_linear = PiccoloOptions(geodesic = false)
+
+# ## Control Initialization
+
+# ### `init_trajectory::Union{NamedTrajectory, Nothing} = nothing`
+# Provide custom initial trajectory instead of automatic initialization.
+
+# traj_custom = initialize_trajectory(...)
+# opts_init = PiccoloOptions(init_trajectory = traj_custom)
+
+# ### `build_trajectory_constraints::Bool = true`
+# Automatically extract constraints from trajectory bounds.
+# Set to `false` if manually managing constraints.
+
+# ## Leakage Suppression (Multilevel Systems)
+
+# ### `leakage_constraint::Bool = false`
+# Add constraint to limit leakage population.
+
+# ### `leakage_constraint_value::Float64 = 1e-3`
+# Maximum allowed leakage: ∑ᵢ |⟨i|ψ⟩|² ≤ leakage_constraint_value
+
+# ### `leakage_cost::Float64 = 1.0`
+# Penalty weight for leakage in objective (soft constraint).
+
+opts_leakage = PiccoloOptions(
+    leakage_constraint = true,
+    leakage_constraint_value = 1e-2,  # 1% max leakage
+    leakage_cost = 1e-1
+)
+
+# Example with embedded operator:
+# sys_transmon = TransmonSystem(levels=5, δ=0.2)
+# op = EmbeddedOperator(:X, sys_transmon)
+# prob_leakage = UnitarySmoothPulseProblem(
+#     sys_transmon, op, T, Δt;
+#     piccolo_options = opts_leakage
+# )
+
+# ## Control Constraints
+
+# ### `complex_control_norm_constraint_name::Union{Symbol, Nothing} = nothing`
+# Apply norm constraint to complex control amplitudes.
+
+# For systems with complex drives (e.g., rotating frame):
+# opts_norm = PiccoloOptions(
+#     complex_control_norm_constraint_name = :u,
+#     complex_control_norm_constraint_radius = 0.2
+# )
+# Enforces: |u_real + i*u_imag| ≤ radius
+
+# ### `complex_control_norm_constraint_radius::Float64 = 1.0`
+# Radius for complex control norm constraint.
+
+# ## Timestep Options
+
+# ### `timesteps_all_equal::Bool = false`
+# Force all timesteps to be equal: Δt[k] = Δt[1] ∀k.
+# Useful for hardware with fixed sampling rates.
+
+opts_equal_dt = PiccoloOptions(timesteps_all_equal = true)
+
+# ## Advanced Dynamics
+
+# ### `rollout_integrator::Symbol = :pade`
+# Integration method for evaluating fidelity.
+# - `:pade`: Padé approximation (default, fast)
+# - `:exp`: Matrix exponential (more accurate)
+
+opts_exp = PiccoloOptions(rollout_integrator = :exp)
+
+# ## Derivative Constraints
+
+# ### `zero_initial_and_final_derivative::Bool = false`
+# Force derivatives to zero at boundaries: u̇[1] = u̇[T] = 0, ü[1] = ü[T] = 0.
+# Creates "smooth ramp" pulses that start and end at zero derivative.
+
+opts_smooth_edges = PiccoloOptions(zero_initial_and_final_derivative = true)
+
+# ## Common Configuration Patterns
+
+# ### Quick and dirty optimization
+opts_quick = PiccoloOptions(
+    verbose = false,
+    geodesic = true
+)
+
+# ### High-fidelity gate
+opts_hifi = PiccoloOptions(
+    verbose = true,
+    bound_state = true,
+    geodesic = true,
+    rollout_integrator = :exp
+)
+
+# ### Multilevel system with leakage suppression
+opts_multilevel = PiccoloOptions(
+    bound_state = true,
+    leakage_constraint = true,
+    leakage_constraint_value = 1e-2,
+    leakage_cost = 1e-1,
+    verbose = true
+)
+
+# ### Smooth pulses for hardware
+opts_hardware = PiccoloOptions(
+    zero_initial_and_final_derivative = true,
+    timesteps_all_equal = true,
+    bound_state = true
+)
+
+# ### Robust optimization
+opts_robust = PiccoloOptions(
+    bound_state = true,
+    geodesic = true,
+    verbose = true
+)
+
+# ## Tips and Tricks
+
+# **When to use `bound_state=true`:**
+# - Optimization struggling to converge
+# - Fidelity stuck below 0.99
+# - Numerical instabilities in state evolution
+# - Working with large Hilbert spaces
+
+# **Leakage vs bounds:**
+# - Leakage *constraint* enforces hard limit (may fail to converge)
+# - Leakage *cost* adds soft penalty (more forgiving)
+# - Use both for best results
+
+# **Geodesic initialization:**
+# - Almost always better than linear
+# - Only disable for debugging or special cases
+# - Particularly important for large T
+
+# **Rollout integrator:**
+# - `:pade` is fast and usually sufficient
+# - `:exp` more accurate for sensitive systems
+# - Both give same result for well-conditioned problems
+
+println("PiccoloOptions configured!")
diff --git a/docs/literate/man/problem_templates_overview.jl b/docs/literate/man/problem_templates_overview.jl
new file mode 100644
index 00000000..dcbec688
--- /dev/null
+++ b/docs/literate/man/problem_templates_overview.jl
@@ -0,0 +1,80 @@
+# # Problem Templates Overview
+
+# QuantumCollocation.jl provides **8 problem templates** that cover common quantum optimal control scenarios. These templates make it easy to set up and solve problems without manually constructing objectives, constraints, and integrators.
+
+# ## Template Comparison
+
+# | Template | State Type | Objective | Time | Use Case |
+# |:---------|:-----------|:----------|:-----|:---------|
+# | [`UnitarySmoothPulseProblem`](@ref) | Unitary | Minimize control effort + infidelity | Fixed | Standard gate synthesis with smooth pulses |
+# | [`UnitaryMinimumTimeProblem`](@ref) | Unitary | Minimize duration | Variable | Fastest gate given fidelity constraint |
+# | [`UnitarySamplingProblem`](@ref) | Unitary | Minimize control effort + infidelity | Fixed | Robust control over multiple systems |
+# | [`UnitaryFreePhase Problem`](@ref) | Unitary | Minimize control effort + infidelity | Fixed | Gate synthesis with free global phase |
+# | [`UnitaryVariationalProblem`](@ref) | Unitary | Minimize control effort + infidelity ± sensitivity | Fixed | Sensitivity/robustness to Hamiltonian terms |
+# | [`QuantumStateSmoothPulseProblem`](@ref) | Ket | Minimize control effort + infidelity | Fixed | State transfer with smooth pulses |
+# | [`QuantumStateMinimumTimeProblem`](@ref) | Ket | Minimize duration | Variable | Fastest state transfer |
+# | [`QuantumStateSamplingProblem`](@ref) | Ket | Minimize control effort + infidelity | Fixed | Robust state transfer over multiple systems |
+
+# ## Key Differences
+
+# ### Unitary vs Ket (Quantum State)
+# - **Unitary problems**: Optimize gate operations (full unitary matrices), commonly used for universal quantum control
+# - **Ket problems**: Optimize state-to-state transfers, useful for initialization and specific state preparation
+
+# ### Smooth Pulse vs Minimum Time
+# - **Smooth Pulse**: Fixed total time `T × Δt`, minimizes control effort with regularization on `u`, `u̇`, `ü`
+# - **Minimum Time**: Variable timesteps `Δt[k]`, minimizes total duration subject to fidelity constraint
+
+# ### Sampling Problems
+# - Solve for a **single control pulse** that works well across **multiple quantum systems**
+# - Useful for robustness against parameter uncertainties or manufacturing variations
+# - Examples: different coupling strengths, detunings, or environmental conditions
+
+# ### Free Phase & Variational
+# - **Free Phase**: Optimizes global phase of target unitary (sometimes easier to reach)
+# - **Variational**: Uses sensitivity analysis to find controls that are robust or sensitive to specific Hamiltonian terms
+
+# ## Quick Selection Guide
+
+# **I want to implement a quantum gate:**
+# - Start simple? → `UnitarySmoothPulseProblem`
+# - Need speed? → `UnitaryMinimumTimeProblem`
+# - Need robustness? → `UnitarySamplingProblem`
+
+# **I want to prepare a quantum state:**
+# - Standard case? → `QuantumStateSmoothPulseProblem`
+# - Speed critical? → `QuantumStateMinimumTimeProblem`
+# - Robust preparation? → `QuantumStateSamplingProblem`
+
+# **I'm tuning my solution:**
+# - Struggling with convergence? → Try `UnitaryFreePhase Problem`
+# - Need parameter sensitivity? → Use `UnitaryVariationalProblem`
+
+# ## Common Parameters
+
+# All templates share these key parameters:
+
+# ```julia
+# prob = UnitarySmoothPulseProblem(
+#     system,           # QuantumSystem defining H(u)
+#     U_goal,           # Target unitary or state
+#     N,                # Number of timesteps
+#     
+#     # Derivative bounds (smoothness)
+#     du_bound = 0.01,            # |u̇| ≤ du_bound
+#     ddu_bound = 0.001,          # |ü| ≤ ddu_bound
+#     
+#     # Regularization weights
+#     R_u = 0.01,                 # Penalize u²
+#     R_du = 0.01,                # Penalize u̇²
+#     R_ddu = 0.01,               # Penalize ü²
+#     
+#     # Initial guess
+#     u_guess = nothing,          # Optional initial controls
+#     
+#     # Advanced options
+#     piccolo_options = PiccoloOptions(...)
+# )
+# ```
+
+# See the individual template pages for parameter details and examples.
diff --git a/docs/literate/man/unitary_problem_templates.jl b/docs/literate/man/unitary_problem_templates.jl
index 7e9ee1ed..41784efb 100644
--- a/docs/literate/man/unitary_problem_templates.jl
+++ b/docs/literate/man/unitary_problem_templates.jl
@@ -19,26 +19,24 @@ instead of driving the system to a target state, the goal is to drive the system
 target unitary operator, `U_goal`.
 
 =#
-T_max = 1.0
-u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-system = QuantumSystem(0.1 * PAULIS.Z, [PAULIS.X, PAULIS.Y], T_max, u_bounds)
-U_goal = GATES.H
+
+system = QuantumSystem(0.1 * PAULIS.Z, [PAULIS.X, PAULIS.Y], 10.0, [1.0, 1.0])
+U_goal = EmbeddedOperator(GATES.H, system)
 N = 51
-Δt = 0.2
 
-prob = UnitarySmoothPulseProblem(system, U_goal, N, Δt);
+prob = UnitarySmoothPulseProblem(system, U_goal, N); 
 
 # _check the fidelity before solving_
 println("Before: ", unitary_rollout_fidelity(prob.trajectory, system))
 
 # _finding an optimal control is as simple as calling `solve!`_
-load_path = joinpath(dirname(Base.active_project()), "data/unitary_problem_templates_89ee72.jld2") # hide
-prob.trajectory = load_traj(load_path) # hide
-nothing # hide
+# load_path = joinpath(dirname(Base.active_project()), "data/unitary_problem_templates_89ee72.jld2") # hide
+# prob.trajectory = load_traj(load_path) # hide
+# nothing # hide
+# solve!(prob, max_iter=100);
 
 #=
 ```julia
-solve!(prob, max_iter=100, verbose=true, print_level=1);
 ```
 
 ```@raw html
@@ -91,13 +89,13 @@ min_prob = UnitaryMinimumTimeProblem(prob, U_goal);
 println("Duration before: ", get_duration(prob.trajectory))
 
 # _solve the minimum time problem_
-load_path = joinpath(dirname(Base.active_project()), "data/unitary_problem_templates_min_time_89ee72.jld2") # hide
-min_prob.trajectory = load_traj(load_path) # hide
-nothing # hide
+# load_path = joinpath(dirname(Base.active_project()), "data/unitary_problem_templates_min_time_89ee72.jld2") # hide
+# min_prob.trajectory = load_traj(load_path) # hide
+# nothing # hide
+# solve!(min_prob, max_iter=100);
 
 #=
 ```julia
-solve!(min_prob, max_iter=100, verbose=true, print_level=1);
 ```
 
 ```@raw html
@@ -134,10 +132,8 @@ This can be useful for exploring robustness, for example.
 =#
 
 # _create a sampling problem_
-T_max = 1.0
-u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-driftless_system = QuantumSystem([PAULIS.X, PAULIS.Y], T_max, u_bounds)
-sampling_prob = UnitarySamplingProblem([system, driftless_system], U_goal, N, Δt);
+driftless_system = QuantumSystem([PAULIS.X, PAULIS.Y], 10.0, [1.0, 1.0])
+sampling_prob = UnitarySamplingProblem([system, driftless_system], U_goal, N);
 
 # _new keys are addded to the trajectory for the new states_
 println(sampling_prob.trajectory.state_names)
@@ -163,9 +159,9 @@ for more details.
 T_max = 1.0
 u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
 H_var = PAULIS.X
-varsys = VariationalQuantumSystem([PAULIS.X, PAULIS.Y], [H_var], T_max, u_bounds);
+varsys = VariationalQuantumSystem([PAULIS.X, PAULIS.Y], [H_var], 10.0, [1.0, 1.0]);
 
 # _create a variational problem that is robust to `PAULIS.X` at the end_
-robprob = UnitaryVariationalProblem(varsys, U_goal, N, Δt, robust_times=[[N]]);
+robprob = UnitaryVariationalProblem(varsys, U_goal, N, robust_times=[[N]]);
 
 # -----
diff --git a/docs/literate/man/working_with_solutions.jl b/docs/literate/man/working_with_solutions.jl
new file mode 100644
index 00000000..0a9acde9
--- /dev/null
+++ b/docs/literate/man/working_with_solutions.jl
@@ -0,0 +1,214 @@
+# # Working with Solutions
+
+# This guide covers how to solve problems, extract data from solutions, evaluate performance, and save/load results.
+
+using QuantumCollocation
+using PiccoloQuantumObjects
+using NamedTrajectories
+
+# ## Solving Problems
+
+# Once you've created a problem template, solving is straightforward with the `solve!` function:
+
+system = QuantumSystem(0.1 * PAULIS.Z, [PAULIS.X, PAULIS.Y], 10.0, [1.0, 1.0])
+U_goal = EmbeddedOperator(GATES.H, system)
+N = 51
+
+prob = UnitarySmoothPulseProblem(system, U_goal, N)
+
+# The `solve!` function accepts several key options:
+
+solve!(prob;
+    max_iter = 100,        # Maximum optimization iterations
+    verbose = true,         # Print convergence information
+    print_level = 1         # Ipopt output level (0-5)
+)
+
+# ### Understanding Convergence
+
+# Ipopt reports several key metrics:
+# - **Objective**: Current cost function value
+# - **inf_pr**: Constraint violation (primal infidelity) - should go to ~0
+# - **inf_du**: Dual infidelity - measure of optimality
+# - **lg(mu)**: Log of barrier parameter
+# - **alpha_du/alpha_pr**: Step sizes
+
+# Successful convergence typically shows `inf_pr` < 1e-6 and status `Optimal Solution Found`.
+
+# ## Extracting Data from Solutions
+
+# ### Accessing Controls
+
+# Control pulses are stored in the trajectory with automatic naming:
+
+u = prob.trajectory.u       # Control values [n_drives × N]
+du = prob.trajectory.du     # First derivatives
+ddu = prob.trajectory.ddu   # Second derivatives
+
+println("Control shape: ", size(u))
+println("Number of drives: ", size(u, 1))
+println("Number of timesteps: ", size(u, 2))
+
+# Access individual drive controls:
+u_drive_1 = u[1, :]   # First drive over time
+
+# ### Accessing States
+
+# For unitary problems:
+Ũ⃗ = prob.trajectory.Ũ⃗  # Vectorized unitary [N² × T]
+
+# The unitary is stored in "isovec" format (vectorized). To get the actual unitary matrix at timestep k:
+using LinearAlgebra
+k = N  # Final timestep
+U_k = iso_vec_to_operator(Ũ⃗[:, k])
+println("Final unitary dimensions: ", size(U_k))
+
+# For ket (state) problems, use:
+# ψ̃ = prob.trajectory.ψ̃  # Vectorized state [2N × T]
+
+# ### Time Grid
+
+# Access timestep information:
+Δt_vec = prob.trajectory.Δt  # Timestep durations
+
+# Calculate total duration:
+duration = get_duration(prob.trajectory)
+println("Total gate time: ", duration, " (arbitrary units)")
+
+# For minimum time problems, timesteps vary:
+# min_prob = UnitaryMinimumTimeProblem(prob, U_goal)
+# solve!(min_prob, max_iter=100)
+# Δt_optimized = min_prob.trajectory.Δt  # Variable timesteps
+
+# ## Evaluating Solutions
+
+# ### Fidelity Calculations
+
+# **Direct fidelity** - Compare final state to goal:
+U_final = iso_vec_to_operator(prob.trajectory.Ũ⃗[:, end])
+fid_direct = unitary_fidelity(U_final, U_goal)
+println("Direct fidelity: ", fid_direct)
+
+# **Rollout fidelity** - Simulate dynamics forward:
+fid_rollout = unitary_rollout_fidelity(prob.trajectory, system)
+println("Rollout fidelity: ", fid_rollout)
+
+# The rollout fidelity is more accurate as it accounts for actual dynamics, while direct fidelity only checks the final point.
+
+# ### For Embedded Operators (Multilevel Systems)
+
+# When working with subspaces (e.g., qubit in transmon):
+# op = EmbeddedOperator(:X, system)
+# prob_embedded = UnitarySmoothPulseProblem(system, op, N)
+# solve!(prob_embedded, max_iter=100)
+# 
+# # Evaluate fidelity only in computational subspace
+# fid_subspace = unitary_rollout_fidelity(
+#     prob_embedded.trajectory, 
+#     system; 
+#     subspace = op.subspace
+# )
+
+# ### Leakage Evaluation
+
+# For multilevel systems, check population in leakage levels:
+# This requires analyzing state populations during evolution
+# See the Multilevel Transmon example for details
+
+# ### Constraint Violations
+
+# Check if solution satisfies all constraints:
+# - Dynamics constraints: Compare rollout vs direct fidelity
+# - Bound constraints: Verify controls within system.drive_bounds
+# - Derivative constraints: Check |du|, |ddu| within bounds
+
+println("Max control value: ", maximum(abs.(u)))
+println("Max control derivative: ", maximum(abs.(du)))
+
+# ## Saving and Loading
+
+# ### Save a Trajectory
+
+# using JLD2
+# save_object("my_solution.jld2", prob.trajectory)
+
+# ### Load and Reuse
+
+# Load trajectory for warm-starting:
+# traj_loaded = load_object("my_solution.jld2")
+# 
+# # Use as initial guess for new problem
+# prob_refined = UnitarySmoothPulseProblem(
+#     system, U_goal, T, Δt;
+#     u_guess = traj_loaded.u,
+#     piccolo_options = PiccoloOptions(verbose=false)
+# )
+
+# ### Save Entire Problem
+
+# To save all problem information including constraints and objectives:
+# save_object("my_problem.jld2", prob)
+
+# ## Post-Processing
+
+# ### Resampling Trajectories
+
+# Change time discretization while preserving control shape:
+# T_new = 101  # More timesteps
+# traj_resampled = resample(prob.trajectory, T_new)
+
+# ### Extracting for Experiments
+
+# Prepare control pulses for hardware:
+
+using PiccoloPlots  # For visualization
+using CairoMakie
+
+# Plot controls
+fig = plot_controls(prob.trajectory)
+# save("controls.png", fig)
+
+# Extract control data for export
+control_data = Dict(
+    "time" => cumsum([0; prob.trajectory.Δt[:]]),
+    "drive_1_real" => u[1, :],
+    "drive_2_real" => u[2, :],
+    "duration" => duration
+)
+
+# # Save to CSV or other format for AWG
+# using CSV, DataFrames
+# df = DataFrame(control_data)
+# CSV.write("pulse_sequence.csv", df)
+
+# ### Pulse Filtering
+
+# Apply smoothing to reduce high-frequency noise:
+# using DSP
+# for i in 1:size(u, 1)
+#     u_filtered[i, :] = filtfilt(responsetype, u[i, :])
+# end
+
+# ## Best Practices
+
+# **Starting a new optimization:**
+# 1. Begin with coarse discretization (small T)
+# 2. Use relaxed bounds and convergence criteria
+# 3. Refine solution incrementally
+# 4. Use previous solution as warm start
+
+# **Debugging poor convergence:**
+# 1. Check `inf_pr` - high values indicate constraint violations
+# 2. Verify system Hamiltonian is correct
+# 3. Try looser derivative bounds (du_bound, ddu_bound)
+# 4. Increase regularization weights (`R_u`, `R_du`, `R_ddu`)
+# 5. Use `piccolo_options.bound_state=true` for better numerics
+
+# **Improving solutions:**
+# 1. Increase T (more timesteps = finer control)
+# 2. Add derivative constraints for smoother pulses
+# 3. Use minimum time optimization for fastest gates
+# 4. Apply leakage constraints for multilevel systems
+# 5. Use sampling problems for robust control
+
+println("Solution evaluation complete!")
diff --git a/docs/make.jl b/docs/make.jl
index 9dfa3058..df5de175 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -4,11 +4,22 @@ using PiccoloDocsTemplate
 pages = [
     "Home" => "index.md",
     "Manual" => [
+        "Problem Templates Overview" => "generated/man/problem_templates_overview.md",
         "Ket Problem Templates" => "generated/man/ket_problem_templates.md",
         "Unitary Problem Templates" => "generated/man/unitary_problem_templates.md",
+        "Working with Solutions" => "generated/man/working_with_solutions.md",
+        "PiccoloOptions Reference" => "generated/man/piccolo_options.md",
+    ],
+    "Customization" => [
+        "Custom Objectives" => "generated/man/custom_objectives.md",
+        "Adding Constraints" => "generated/man/adding_constraints.md",
+        "Initial Trajectories" => "generated/man/initial_trajectories.md",
     ],
     "Examples" => [
+        "Single Qubit Gate" => "generated/examples/single_qubit_gate.md",
         "Two Qubit Gates" => "generated/examples/two_qubit_gates.md",
+        "Minimum Time Optimization" => "generated/examples/minimum_time_problem.md",
+        "Robust Control" => "generated/examples/robust_control.md",
         "Multilevel Transmon" => "generated/examples/multilevel_transmon.md",
     ],
     "Library" => "lib.md",
diff --git a/docs/quantum_trajectory_architecture.md b/docs/quantum_trajectory_architecture.md
new file mode 100644
index 00000000..4382919b
--- /dev/null
+++ b/docs/quantum_trajectory_architecture.md
@@ -0,0 +1,130 @@
+# Quantum Trajectory Architecture
+
+## Overview
+
+We've introduced `AbstractQuantumTrajectory` types that wrap `NamedTrajectory` with quantum-specific metadata. This enables type-based dispatch for problem templates and cleaner code organization.
+
+## Type Hierarchy
+
+```julia
+abstract type AbstractQuantumTrajectory end
+
+struct UnitaryTrajectory <: AbstractQuantumTrajectory
+    trajectory::NamedTrajectory
+    system::QuantumSystem
+    state_name::Symbol    # :Ũ⃗
+    control_name::Symbol  # :u
+    goal::AbstractPiccoloOperator
+end
+
+struct KetTrajectory <: AbstractQuantumTrajectory
+    trajectory::NamedTrajectory
+    system::QuantumSystem
+    state_name::Symbol    # :ψ̃
+    control_name::Symbol  # :u
+    goals::Vector{<:AbstractVector{ComplexF64}}
+end
+
+struct DensityTrajectory <: AbstractQuantumTrajectory
+    trajectory::NamedTrajectory
+    system::OpenQuantumSystem
+    state_name::Symbol    # :ρ⃗̃
+    control_name::Symbol  # :u
+    goal::AbstractMatrix
+end
+```
+
+## Usage Example
+
+```julia
+using QuantumCollocation
+using PiccoloQuantumObjects
+
+# Step 1: Create quantum system
+sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+
+# Step 2: Create quantum trajectory (returns wrapped type)
+qtraj = unitary_trajectory(sys, GATES[:H], 51; store_times=true)
+# Returns: UnitaryTrajectory
+
+# Access underlying data
+qtraj[:Ũ⃗]  # state data
+qtraj[:u]   # control data
+
+# Access metadata
+get_system(qtraj)       # the QuantumSystem
+get_goal(qtraj)         # target operator
+get_state_name(qtraj)   # :Ũ⃗
+get_control_name(qtraj) # :u
+
+# Convenience accessors
+state(qtraj)     # qtraj[:Ũ⃗]
+controls(qtraj)  # qtraj[:u]
+```
+
+## Benefits
+
+1. **Type Safety**: Trajectory type determines which problem templates and integrators apply
+2. **Cleaner API**: Metadata travels with trajectory, no need to pass separately
+3. **Delegation**: Common operations delegate to underlying `NamedTrajectory`
+4. **Multiple Dispatch**: Problem templates can dispatch on quantum trajectory type
+
+## Next Steps
+
+### 1. Integrator Partial Construction
+
+Create integrators that can be configured without being bound to a trajectory:
+
+```julia
+# Configure integrator with kwargs
+int_spec = UnitarySplineIntegrator(spline_order=3, tol=1e-8)
+
+# Later, bind to trajectory
+int_bound = int_spec(qtraj)  # or int_spec(sys, traj, :Ũ⃗, :u)
+```
+
+### 2. SmoothPulseProblem Template
+
+Single problem template that works for all quantum trajectory types:
+
+```julia
+# Dispatches on UnitaryTrajectory
+prob = SmoothPulseProblem(
+    qtraj::UnitaryTrajectory;
+    integrator=UnitarySplineIntegrator(spline_order=3, tol=1e-8),
+    Q=100.0,
+    R=1e-2
+)
+
+# Internally:
+# - Adds derivatives to trajectory
+# - Creates appropriate objective (UnitaryInfidelityObjective)
+# - Instantiates integrators
+# - Assembles DirectTrajOptProblem
+```
+
+### 3. Trajectory Augmentation
+
+Add derivatives within problem template:
+
+```julia
+function SmoothPulseProblem(qtraj::AbstractQuantumTrajectory; kwargs...)
+    # Add control derivatives
+    traj_smooth = add_derivatives(qtraj.trajectory, :u, 2; bounds=...)
+    
+    # Create objectives, integrators, etc.
+    ...
+end
+```
+
+## Implementation Status
+
+✅ Created `AbstractQuantumTrajectory` type hierarchy  
+✅ Updated trajectory creators to return wrapped types  
+✅ Implemented delegation to underlying `NamedTrajectory`  
+✅ Added accessor functions  
+✅ Updated tests
+
+⏳ Next: Integrator partial construction pattern  
+⏳ Next: Refactor problem templates to use quantum trajectories  
+⏳ Next: Implement `SmoothPulseProblem` with type dispatch
diff --git a/docs/src/index.md b/docs/src/index.md
index aaa32321..45f01219 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -1,62 +1,135 @@
 # QuantumCollocation.jl
 
-**QuantumCollocation.jl** sets up and solves *quantum control problems* as nonlinear programs (NLPs). In this context, a generic quantum control problem looks like
+**QuantumCollocation.jl** is a Julia package for solving quantum optimal control problems using direct collocation methods. It transforms continuous-time quantum control into finite-dimensional nonlinear programs (NLPs) solved efficiently with [Ipopt](https://github.com/jump-dev/Ipopt.jl).
+
+## Quick Start
+
+### Installation
+
+```julia
+using Pkg
+Pkg.add("QuantumCollocation")
+```
+
+Or from the Julia REPL, press `]` to enter package mode:
+```
+pkg> add QuantumCollocation
+```
+
+### 30-Second Example
+
+Here's a complete example optimizing a Hadamard gate on a qubit:
+
+```julia
+using QuantumCollocation
+using PiccoloQuantumObjects
+
+# Define system: drift + 2 control Hamiltonians
+H_drift = 0.1 * PAULIS.Z
+H_drives = [PAULIS.X, PAULIS.Y]
+sys = QuantumSystem(H_drift, H_drives, 10.0, [1.0, 1.0])
+
+# Set up problem: system, target gate, timesteps
+U_goal = GATES.H
+N = 51
+prob = UnitarySmoothPulseProblem(sys, U_goal, N)
+
+# Solve!
+solve!(prob; max_iter=100)
+
+# Check result
+println("Fidelity: ", unitary_rollout_fidelity(prob.trajectory, sys))
+```
+
+That's it! You've optimized control pulses for a quantum gate.
+
+## What Can QuantumCollocation Do?
+
+- **Unitary gate optimization** - Find pulses to implement quantum gates
+- **State transfer** - Drive quantum states to target states
+- **Minimum time control** - Optimize gate duration
+- **Robust control** - Account for system uncertainties
+- **Multilevel systems** - Handle transmons, bosonic codes, etc.
+- **Leakage suppression** - Constrain populations in unwanted levels
+- **Custom constraints** - Add your own physics constraints
+
+## Overview
+
+**QuantumCollocation.jl** sets up and solves *quantum control problems* as nonlinear programs (NLPs). A generic quantum control problem looks like:
+
 ```math
 \begin{aligned}
     \arg \min_{\mathbf{Z}}\quad & J(\mathbf{Z}) \\
-    \nonumber \text{s.t.}\qquad & \mathbf{f}(\mathbf{Z}) = 0 \\
-    \nonumber & \mathbf{g}(\mathbf{Z}) \le 0  
+    \text{s.t.}\qquad & \mathbf{f}(\mathbf{Z}) = 0 \\
+    & \mathbf{g}(\mathbf{Z}) \le 0
 \end{aligned}
 ```
-where $\mathbf{Z}$ is a trajectory  containing states and controls, from [NamedTrajectories.jl](https://github.com/harmoniqs/NamedTrajectories.jl).
 
------
+where $\mathbf{Z}$ is a trajectory containing states and controls from [NamedTrajectories.jl](https://github.com/harmoniqs/NamedTrajectories.jl).
 
-We provide a number of **problem templates** for making it simple and easy to set up and solve 
-certain types of quantum optimal control problems. These templates all construct a 
-`DirectTrajOptProblem` object from [DirectTrajOpt.jl](https://github.com/harmoniqs/DirectTrajOpt.jl), which stores all the parts of the optimal control problem.
+We provide **problem templates** for common quantum control tasks. These templates construct a `DirectTrajOptProblem` from [DirectTrajOpt.jl](https://github.com/harmoniqs/DirectTrajOpt.jl) with appropriate objectives, constraints, and dynamics.
 
------
+We provide **problem templates** for common quantum control tasks. These templates construct a `DirectTrajOptProblem` from [DirectTrajOpt.jl](https://github.com/harmoniqs/DirectTrajOpt.jl) with appropriate objectives, constraints, and dynamics.
 
-### Get started
+## Problem Templates
 
-The problem templates are broken down by the state variable of the problem being solved.
+Problem templates are organized by the type of quantum system being controlled:
 
-Ket Problem Templates:
-- [Quantum State Smooth Pulse Problem](@ref)
-- [Quantum State Minimum Time Problem](@ref)
-- [Quantum State Sampling Problem](@ref)
+### Unitary (Gate) Templates
+- [`UnitarySmoothPulseProblem`](@ref) - Optimize smooth pulses for unitary gates
+- [`UnitaryMinimumTimeProblem`](@ref) - Minimize gate duration
+- [`UnitarySamplingProblem`](@ref) - Robust control over system variations
+- [`UnitaryFreePhaseProblem`](@ref) - Optimize up to global phase
+- [`UnitaryVariationalProblem`](@ref) - Variational quantum optimization
 
-Unitary Problem Templates:
-- [Unitary Smooth Pulse Problem](@ref)
-- [Unitary Minimum Time Problem](@ref)
-- [Unitary Sampling Problem](@ref)
+### Quantum State Templates
+- [`QuantumStateSmoothPulseProblem`](@ref) - Drive states with smooth pulses
+- [`QuantumStateMinimumTimeProblem`](@ref) - Minimize state transfer time
+- [`QuantumStateSamplingProblem`](@ref) - Robust state transfer
 
-### Background
+See the [Problem Templates Overview](@ref) for a detailed comparison and selection guide.
 
-*Problem Templates* are reusable design patterns for setting up and solving common quantum control problems. 
+See the [Problem Templates Overview](@ref) for a detailed comparison and selection guide.
+
+## How It Works
+
+### Direct Collocation
+
+QuantumCollocation uses *direct collocation* - discretizing continuous-time dynamics into constraints at discrete time points (knot points). For example, a smooth pulse problem for a unitary gate:
 
-For example, a *UnitarySmoothPulseProblem* is tasked with generating a *pulse* sequence $a_{1:T-1}$ in orderd to minimize infidelity, subject to constraints from the Schroedinger equation,
-```math
-    \begin{aligned}
-        \arg \min_{\mathbf{Z}}\quad & |1 - \mathcal{F}(U_T, U_\text{goal})|  \\
-        \nonumber \text{s.t.}
-        \qquad & U_{t+1} = \exp\{- i H(a_t) \Delta t_t \} U_t, \quad \forall\, t \\
-    \end{aligned}
-```
-while a *UnitaryMinimumTimeProblem* minimizes time and constrains fidelity,
 ```math
-    \begin{aligned}
-        \arg \min_{\mathbf{Z}}\quad & \sum_{t=1}^T \Delta t_t \\
-        \qquad & U_{t+1} = \exp\{- i H(a_t) \Delta t_t \} U_t, \quad \forall\, t \\
-        \nonumber & \mathcal{F}(U_T, U_\text{goal}) \ge 0.9999
-    \end{aligned}
+\begin{aligned}
+    \arg \min_{\mathbf{Z}}\quad & |1 - \mathcal{F}(U_N, U_\text{goal})|  \\
+    \text{s.t.}\qquad & U_{k+1} = \exp\{- i H(u_k) \Delta t_k \} U_k, \quad \forall\, k \\
+    & u_{k+1} = u_k + \dot{u}_k \Delta t_k \\
+    & \dot{u}_{k+1} = \dot{u}_k + \ddot{u}_k \Delta t_k \\
+    & |u_k| \le u_\text{max}, \quad |\ddot{u}_k| \le \ddot{u}_\text{max}
+\end{aligned}
 ```
 
------
+The dynamics between knot points $(U_k, u_k)$ and $(U_{k+1}, u_{k+1})$ become nonlinear equality constraints. States and controls are free variables optimized by the NLP solver.
+
+### Key Features
+
+- **Efficient gradients** - Sparse Jacobians and Hessians via automatic differentiation
+- **Flexible constraints** - Add custom physics, leakage suppression, robustness
+- **Multiple integrators** - Exponential, Pade, time-dependent dynamics
+- **Extensible** - Easy to add new objectives, constraints, and problem templates
+
+## Next Steps
+
+- 📚 [Problem Templates Overview](@ref) - Choose the right template for your problem
+- 🎯 [Working with Solutions](@ref) - Extract results, evaluate fidelity, save data
+- ⚙️ [PiccoloOptions Reference](@ref) - Configure solver options and constraints
+- 💡 [Examples](@ref) - See complete examples from single qubits to multilevel systems
+
+## Related Packages
 
-In each case, the dynamics between *knot points* $(U_t, a_t)$ and $(U_{t+1}, a_{t+1})$ are enforced as constraints on the states, which are free variables in the solver; this optimization framework is called *direct trajectory optimization*. 
+QuantumCollocation.jl is part of the [Piccolo ecosystem](https://github.com/harmoniqs/Piccolo.jl):
 
------
+- [**NamedTrajectories.jl**](https://github.com/harmoniqs/NamedTrajectories.jl) - Trajectory data structures
+- [**DirectTrajOpt.jl**](https://github.com/harmoniqs/DirectTrajOpt.jl) - Direct trajectory optimization framework
+- [**PiccoloQuantumObjects.jl**](https://github.com/harmoniqs/PiccoloQuantumObjects.jl) - Quantum operators and systems
+- [**PiccoloPlots.jl**](https://github.com/harmoniqs/PiccoloPlots.jl) - Visualization tools
 
 Problem templates give the user the ability to add other constraints and objective functions to this problem and solve it efficiently using [Ipopt.jl](https://github.com/jump-dev/Ipopt.jl) and [MathOptInterface.jl](https://github.com/jump-dev/MathOptInterface.jl) under the hood (support for additional backends coming soon!).
diff --git a/src/QuantumCollocation.jl b/src/QuantumCollocation.jl
index 700e175d..a19e3714 100644
--- a/src/QuantumCollocation.jl
+++ b/src/QuantumCollocation.jl
@@ -2,13 +2,14 @@ module QuantumCollocation
 
 using Reexport
 
-@reexport using DirectTrajOpt
+using DirectTrajOpt
+using PiccoloQuantumObjects
 
 include("piccolo_options.jl")
 @reexport using .Options
 
-include("trajectory_initialization.jl")
-@reexport using .TrajectoryInitialization
+include("quantum_control_problem.jl")
+@reexport using .QuantumControlProblems
 
 include("quantum_objectives.jl")
 @reexport using .QuantumObjectives
@@ -22,7 +23,4 @@ include("quantum_integrators.jl")
 include("problem_templates/_problem_templates.jl")
 @reexport using .ProblemTemplates
 
-include("quantum_system_templates/_quantum_system_templates.jl")
-@reexport using .QuantumSystemTemplates
-
 end
diff --git a/src/problem_templates/_problem_templates.jl b/src/problem_templates/_problem_templates.jl
index 6656b57f..8961ea47 100644
--- a/src/problem_templates/_problem_templates.jl
+++ b/src/problem_templates/_problem_templates.jl
@@ -1,38 +1,35 @@
 module ProblemTemplates
 
-using ..TrajectoryInitialization
 using ..QuantumObjectives
 using ..QuantumConstraints
 using ..QuantumIntegrators
+using ..QuantumControlProblems: QuantumControlProblem
 using ..Options
 
 using TrajectoryIndexingUtils
 using NamedTrajectories
 using DirectTrajOpt
 using PiccoloQuantumObjects
+using PiccoloQuantumObjects: build_sampling_trajectory, build_ensemble_trajectory_from_trajectories, 
+    get_ensemble_state_names, get_weights, update_base_trajectory, SamplingTrajectory, EnsembleTrajectory
 
 using ExponentialAction
 using LinearAlgebra
 using SparseArrays
 using TestItems
 
-include("unitary_smooth_pulse_problem.jl")
-include("unitary_variational_problem.jl")
-include("unitary_minimum_time_problem.jl")
-include("unitary_sampling_problem.jl")
-include("unitary_free_phase_problem.jl")
-
-include("quantum_state_smooth_pulse_problem.jl")
-include("quantum_state_minimum_time_problem.jl")
-include("quantum_state_sampling_problem.jl")
+const ⊗ = kron
 
+include("smooth_pulse_problem.jl")
+include("minimum_time_problem.jl")
+include("sampling_problem.jl")
 
 function apply_piccolo_options!(
     piccolo_options::PiccoloOptions,
     constraints::AbstractVector{<:AbstractConstraint},
     traj::NamedTrajectory;
-    state_names::Union{Nothing, Symbol, AbstractVector{Symbol}}=nothing,
-    state_leakage_indices::Union{Nothing, AbstractVector{Int}, AbstractVector{<:AbstractVector{Int}}}=nothing,
+    state_names::Union{Nothing,Symbol,AbstractVector{Symbol}}=nothing,
+    state_leakage_indices::Union{Nothing,AbstractVector{Int},AbstractVector{<:AbstractVector{Int}}}=nothing,
 )
     J = NullObjective(traj)
 
@@ -63,7 +60,7 @@ function apply_piccolo_options!(
         end
         push!(
             constraints,
-            TimeStepsAllEqualConstraint(traj)
+            TimeStepsAllEqualConstraint()
         )
     end
 
diff --git a/src/problem_templates/minimum_time_problem.jl b/src/problem_templates/minimum_time_problem.jl
new file mode 100644
index 00000000..e5e5a58a
--- /dev/null
+++ b/src/problem_templates/minimum_time_problem.jl
@@ -0,0 +1,659 @@
+export MinimumTimeProblem
+
+@doc raw"""
+    MinimumTimeProblem(qcp::QuantumControlProblem; kwargs...)
+
+Convert an existing quantum control problem to minimum-time optimization.
+
+**IMPORTANT**: This function requires an existing `QuantumControlProblem` (e.g., from `SmoothPulseProblem`).
+It cannot be created directly from a quantum trajectory. The workflow is:
+1. Create base problem with `SmoothPulseProblem` (or similar)
+2. Solve base problem to get feasible solution
+3. Convert to minimum-time with `MinimumTimeProblem`
+
+This ensures the problem starts from a good initialization and maintains solution quality
+through the final fidelity constraint.
+
+# Type Dispatch
+Automatically handles different quantum trajectory types through the type parameter:
+- `QuantumControlProblem{UnitaryTrajectory}` → Uses `FinalUnitaryFidelityConstraint`
+- `QuantumControlProblem{KetTrajectory}` → Uses `FinalKetFidelityConstraint`
+- `QuantumControlProblem{DensityTrajectory}` → Not yet implemented
+
+The optimization problem is:
+
+```math
+\begin{aligned}
+\underset{\vec{\tilde{q}}, u, \Delta t}{\text{minimize}} & \quad
+J_{\text{original}}(\vec{\tilde{q}}, u) + D \sum_t \Delta t_t \\
+\text{ subject to } & \quad \text{original dynamics \& constraints} \\
+& F_{\text{final}} \geq F_{\text{threshold}} \\
+& \quad \Delta t_{\text{min}} \leq \Delta t_t \leq \Delta t_{\text{max}} \\
+\end{aligned}
+```
+
+where q represents the quantum state (unitary, ket, or density matrix).
+
+# Arguments
+- `qcp::QuantumControlProblem`: Existing quantum control problem to convert
+
+# Keyword Arguments
+- `final_fidelity::Float64=0.99`: Minimum fidelity constraint at final time
+- `D::Float64=100.0`: Weight on minimum-time objective ∑Δt
+- `piccolo_options::PiccoloOptions=PiccoloOptions()`: Piccolo solver options
+
+# Returns
+- `QuantumControlProblem`: New problem with minimum-time objective and fidelity constraint
+
+# Examples
+```julia
+# Standard workflow
+sys = QuantumSystem(H_drift, H_drives, T, drive_bounds)
+qtraj = UnitaryTrajectory(sys, U_goal, N; Δt_bounds=(0.01, 0.5))
+
+# Step 1: Create and solve base smooth pulse problem
+qcp_smooth = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+solve!(qcp_smooth; max_iter=100)
+
+# Step 2: Convert to minimum-time
+qcp_mintime = MinimumTimeProblem(qcp_smooth; final_fidelity=0.99, D=100.0)
+solve!(qcp_mintime; max_iter=100)
+
+# Compare durations
+duration_before = sum(get_timesteps(get_trajectory(qcp_smooth)))
+duration_after = sum(get_timesteps(get_trajectory(qcp_mintime)))
+@assert duration_after <= duration_before
+
+# Nested transformations also work
+qcp_final = MinimumTimeProblem(
+    RobustnessProblem(qcp_smooth);  # Future feature
+    final_fidelity=0.95
+)
+```
+
+# Convenience Constructors
+
+You can also update the goal when creating minimum-time problem:
+
+```julia
+# Different goal for minimum-time optimization
+qcp_mintime = MinimumTimeProblem(qcp_smooth; goal=U_goal_new, final_fidelity=0.98)
+```
+"""
+function MinimumTimeProblem(
+    qcp::QuantumControlProblem{QT};
+    goal::Union{Nothing,AbstractPiccoloOperator,AbstractVector}=nothing,
+    final_fidelity::Float64=0.99,
+    D::Float64=100.0,
+    piccolo_options::PiccoloOptions=PiccoloOptions()
+) where {QT<:AbstractQuantumTrajectory}
+
+    if piccolo_options.verbose
+        println("    constructing MinimumTimeProblem from QuantumControlProblem{$QT}...")
+        println("\tfinal fidelity constraint: $(final_fidelity)")
+        println("\tminimum-time weight D: $(D)")
+    end
+
+    # Copy trajectory and constraints from original problem
+    traj = deepcopy(qcp.prob.trajectory)
+    constraints = deepcopy(qcp.prob.constraints)
+
+    # Add minimum-time objective to existing objective
+    J = qcp.prob.objective + MinimumTimeObjective(traj, D=D)
+
+    # Use updated goal if provided, otherwise use original
+    qtraj_for_constraint = if isnothing(goal)
+        qcp.qtraj
+    else
+        # Create new quantum trajectory with updated goal
+        _update_goal(qcp.qtraj, goal)
+    end
+
+    # Add final fidelity constraint - dispatches on QT type parameter!
+    fidelity_constraint = _final_fidelity_constraint(
+        qtraj_for_constraint,
+        final_fidelity,
+        traj
+    )
+
+    # Handle single constraint or multiple constraints
+    if fidelity_constraint isa AbstractVector
+        append!(constraints, fidelity_constraint)
+    else
+        push!(constraints, fidelity_constraint)
+    end
+
+    # Create new optimization problem with same integrators
+    new_prob = DirectTrajOptProblem(
+        traj,
+        J,
+        qcp.prob.integrators,
+        constraints
+    )
+
+    # Return new QuantumControlProblem with potentially updated qtraj
+    return QuantumControlProblem(qtraj_for_constraint, new_prob)
+end
+
+# ============================================================================= #
+# Type-specific helper functions
+# ============================================================================= #
+
+# Helper to update goal in quantum trajectory (convenience constructor support)
+function _update_goal(qtraj::UnitaryTrajectory, new_goal::AbstractPiccoloOperator)
+    # Create new trajectory with updated goal, preserving Δt_bounds
+    traj = PiccoloQuantumObjects.get_trajectory(qtraj)
+    Δt_bounds = if haskey(traj.bounds, :Δt)
+        # Extract scalar bounds from vector bounds (Δt is 1D)
+        lb, ub = traj.bounds.Δt
+        (lb[1], ub[1])
+    else
+        nothing
+    end
+
+    return UnitaryTrajectory(
+        PiccoloQuantumObjects.get_system(qtraj),
+        new_goal,
+        traj.N;
+        Δt_bounds=Δt_bounds,
+    )
+end
+
+function _update_goal(qtraj::KetTrajectory, new_goal::AbstractVector{<:Number})
+    # Keep initial state, update goal
+    traj = PiccoloQuantumObjects.get_trajectory(qtraj)
+    Δt_bounds = if haskey(traj.bounds, :Δt)
+        # Extract scalar bounds from vector bounds (Δt is 1D)
+        lb, ub = traj.bounds.Δt
+        (lb[1], ub[1])
+    else
+        nothing
+    end
+
+    return KetTrajectory(
+        PiccoloQuantumObjects.get_system(qtraj),
+        get_state(qtraj),  # Keep initial state
+        new_goal,          # New goal state
+        traj.N;
+        Δt_bounds=Δt_bounds,
+    )
+end
+
+function _update_goal(qtraj::DensityTrajectory, new_goal::AbstractMatrix)
+    traj = PiccoloQuantumObjects.get_trajectory(qtraj)
+    Δt_bounds = if haskey(traj.bounds, :Δt)
+        # Extract scalar bounds from vector bounds (Δt is 1D)
+        lb, ub = traj.bounds.Δt
+        (lb[1], ub[1])
+    else
+        nothing
+    end
+
+    return DensityTrajectory(
+        PiccoloQuantumObjects.get_system(qtraj),
+        get_state(qtraj),  # Keep initial state
+        new_goal,      # New goal state
+        traj.N;
+        Δt_bounds=Δt_bounds,
+    )
+end
+
+# Fidelity constraint functions - dispatch on quantum trajectory type
+
+function _final_fidelity_constraint(
+    qtraj::UnitaryTrajectory,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    U_goal = PiccoloQuantumObjects.get_goal(qtraj)
+    state_sym = PiccoloQuantumObjects.get_state_name(qtraj)
+    return FinalUnitaryFidelityConstraint(U_goal, state_sym, final_fidelity, traj)
+end
+
+function _final_fidelity_constraint(
+    qtraj::KetTrajectory,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    ψ_goal = PiccoloQuantumObjects.get_goal(qtraj)
+    state_sym = PiccoloQuantumObjects.get_state_name(qtraj)
+    return FinalKetFidelityConstraint(ψ_goal, state_sym, final_fidelity, traj)
+end
+
+function _final_fidelity_constraint(
+    qtraj::DensityTrajectory,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    # TODO: Implement density matrix fidelity constraint when available
+    throw(ArgumentError("Final fidelity constraint for DensityTrajectory not yet implemented"))
+end
+
+# ============================================================================= #
+# EnsembleTrajectory Support
+# ============================================================================= #
+
+"""
+    _final_fidelity_constraint(qtraj::EnsembleTrajectory, final_fidelity, traj)
+
+Create fidelity constraints for each trajectory in an EnsembleTrajectory.
+Returns a vector of constraints, one per ensemble member.
+"""
+function _final_fidelity_constraint(
+    qtraj::EnsembleTrajectory,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    state_names = get_ensemble_state_names(qtraj)
+    goals = PiccoloQuantumObjects.get_goal(qtraj)
+    
+    # Create one fidelity constraint per ensemble member
+    constraints = [
+        _ensemble_fidelity_constraint(qtraj.trajectories[1], goal, name, final_fidelity, traj)
+        for (goal, name) in zip(goals, state_names)
+    ]
+    
+    return constraints
+end
+
+# Dispatch on base trajectory type for ensemble fidelity constraint
+function _ensemble_fidelity_constraint(
+    base_qtraj::UnitaryTrajectory,
+    goal::AbstractPiccoloOperator,
+    state_sym::Symbol,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    return FinalUnitaryFidelityConstraint(goal, state_sym, final_fidelity, traj)
+end
+
+function _ensemble_fidelity_constraint(
+    base_qtraj::KetTrajectory,
+    goal::AbstractVector,
+    state_sym::Symbol,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    return FinalKetFidelityConstraint(goal, state_sym, final_fidelity, traj)
+end
+
+function _ensemble_fidelity_constraint(
+    base_qtraj::DensityTrajectory,
+    goal::AbstractMatrix,
+    state_sym::Symbol,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    throw(ArgumentError("Final fidelity constraint for DensityTrajectory ensemble not yet implemented"))
+end
+
+# Update goal for EnsembleTrajectory is not typically needed since goals are embedded
+# in the individual trajectories. But we provide a fallback that errors clearly.
+function _update_goal(qtraj::EnsembleTrajectory, new_goal)
+    throw(ArgumentError(
+        "Updating goals for EnsembleTrajectory is not directly supported. " *
+        "Create a new EnsembleTrajectory with the desired goals instead."
+    ))
+end
+
+# ============================================================================= #
+# Tests
+# ============================================================================= #
+
+@testitem "MinimumTimeProblem from SmoothPulseProblem (Unitary)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+
+    # Create and solve smooth pulse problem
+    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    qtraj = UnitaryTrajectory(sys, GATES[:H], 50; Δt_bounds=(0.01, 0.5))
+    qcp_smooth = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+
+    solve!(qcp_smooth; max_iter=50, verbose=false, print_level=1)
+    duration_before = sum(get_timesteps(get_trajectory(qcp_smooth)))
+
+    # Convert to minimum-time problem
+    qcp_mintime = MinimumTimeProblem(qcp_smooth; final_fidelity=0.95, D=100.0)
+
+    @test qcp_mintime isa QuantumControlProblem
+    @test qcp_mintime isa QuantumControlProblem{UnitaryTrajectory}
+    @test haskey(get_trajectory(qcp_mintime).components, :du)
+    @test haskey(get_trajectory(qcp_mintime).components, :ddu)
+
+    # Test accessors
+    @test get_system(qcp_mintime) === sys
+    @test get_goal(qcp_mintime) === GATES[:H]
+
+    # Solve minimum-time problem
+    solve!(qcp_mintime; max_iter=50, verbose=false, print_level=1)
+    duration_after = sum(get_timesteps(get_trajectory(qcp_mintime)))
+
+    # Duration should decrease (or stay same if already optimal)
+    @test duration_after <= duration_before
+end
+
+@testitem "MinimumTimeProblem from SmoothPulseProblem (Ket)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+    qtraj = KetTrajectory(sys, ψ_init, ψ_goal, 30; Δt_bounds=(0.01, 0.5))
+
+    # Create smooth pulse problem
+    qcp_smooth = SmoothPulseProblem(qtraj; Q=50.0, R=1e-3)
+    solve!(qcp_smooth; max_iter=10, verbose=false, print_level=1)
+
+    # Convert to minimum-time
+    qcp_mintime = MinimumTimeProblem(qcp_smooth; final_fidelity=0.90, D=50.0)
+
+    @test qcp_mintime isa QuantumControlProblem{KetTrajectory}
+    @test haskey(get_trajectory(qcp_mintime).components, :du)
+
+    # Test problem solve
+    solve!(qcp_mintime; max_iter=10, print_level=1, verbose=false)
+end
+
+@testitem "MinimumTimeProblem with updated goal" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    qtraj = UnitaryTrajectory(sys, GATES[:H], 20; Δt_bounds=(0.01, 0.5))
+
+    qcp_smooth = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+    solve!(qcp_smooth; max_iter=5, verbose=false, print_level=1)
+
+    # Create minimum-time with different goal
+    qcp_mintime = MinimumTimeProblem(
+        qcp_smooth;
+        goal=GATES[:X],  # Different goal!
+        final_fidelity=0.95,
+        D=100.0
+    )
+
+    @test qcp_mintime isa QuantumControlProblem
+    @test get_goal(qcp_mintime) === GATES[:X]  # Goal should be updated
+    @test get_goal(qcp_mintime) !== get_goal(qcp_smooth)  # Different from original
+end
+
+@testitem "MinimumTimeProblem type dispatch" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Test that type parameter is correct for different trajectory types
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+
+    # Unitary
+    qtraj_u = UnitaryTrajectory(sys, GATES[:H], 10)
+    qcp_u = SmoothPulseProblem(qtraj_u)
+    qcp_mintime_u = MinimumTimeProblem(qcp_u)
+    @test qcp_mintime_u isa QuantumControlProblem{UnitaryTrajectory}
+
+    # Ket
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+    qtraj_k = KetTrajectory(sys, ψ_init, ψ_goal, 10)
+    qcp_k = SmoothPulseProblem(qtraj_k)
+    qcp_mintime_k = MinimumTimeProblem(qcp_k)
+    @test qcp_mintime_k isa QuantumControlProblem{KetTrajectory}
+end
+
+@testitem "MinimumTimeProblem with SamplingTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+
+    # Robust minimum-time gate optimization
+    sys_nominal = QuantumSystem(0.1 * GATES[:Z], [GATES[:X]], 1.0, [1.0])
+    sys_perturbed = QuantumSystem(0.11 * GATES[:Z], [GATES[:X]], 1.0, [1.0])
+
+    qtraj = UnitaryTrajectory(sys_nominal, GATES[:X], 30; Δt_bounds=(0.01, 0.5))
+    qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+
+    # Create sampling problem
+    sampling_prob = SamplingProblem(qcp, [sys_nominal, sys_perturbed]; Q=100.0)
+    solve!(sampling_prob; max_iter=20, verbose=false, print_level=1)
+
+    duration_before = sum(get_timesteps(get_trajectory(sampling_prob)))
+
+    # Convert to minimum-time
+    mintime_prob = MinimumTimeProblem(sampling_prob; final_fidelity=0.90, D=50.0)
+
+    @test mintime_prob isa QuantumControlProblem{<:SamplingTrajectory}
+    @test mintime_prob.qtraj isa SamplingTrajectory
+
+    # Should have fidelity constraints for each sample
+    # (one per system in the sampling ensemble)
+
+    # Solve minimum-time
+    solve!(mintime_prob; max_iter=50, verbose=false, print_level=1)
+
+    duration_after = sum(get_timesteps(get_trajectory(mintime_prob)))
+    @test duration_after <= duration_before * 1.1  # Allow small tolerance
+end
+
+@testitem "MinimumTimeProblem with SamplingTrajectory (Ket)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Robust minimum-time state transfer
+    sys_nominal = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    sys_perturbed = QuantumSystem(1.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+    qtraj = KetTrajectory(sys_nominal, ψ_init, ψ_goal, 25; Δt_bounds=(0.01, 0.5))
+
+    qcp = SmoothPulseProblem(qtraj; Q=50.0, R=1e-3)
+
+    # Create sampling problem
+    sampling_prob = SamplingProblem(qcp, [sys_nominal, sys_perturbed]; Q=50.0)
+    solve!(sampling_prob; max_iter=15, verbose=false, print_level=1)
+
+    # Convert to minimum-time
+    mintime_prob = MinimumTimeProblem(sampling_prob; final_fidelity=0.85, D=30.0)
+
+    @test mintime_prob isa QuantumControlProblem{<:SamplingTrajectory}
+
+    # Solve
+    solve!(mintime_prob; max_iter=15, verbose=false, print_level=1)
+end
+@testitem "MinimumTimeProblem with EnsembleTrajectory (Ket)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+    using LinearAlgebra
+
+    # Multi-state minimum-time optimization (X gate via state transfer)
+    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    
+    # Create ensemble of ket trajectories for X gate
+    ψ0 = ComplexF64[1.0, 0.0]
+    ψ1 = ComplexF64[0.0, 1.0]
+    
+    qtraj1 = KetTrajectory(sys, ψ0, ψ1, 30; Δt_bounds=(0.01, 0.5))  # |0⟩ → |1⟩
+    qtraj2 = KetTrajectory(sys, ψ1, ψ0, 30; Δt_bounds=(0.01, 0.5))  # |1⟩ → |0⟩
+    
+    ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+    
+    # Create and solve smooth pulse problem
+    qcp_smooth = SmoothPulseProblem(ensemble_qtraj; Q=100.0, R=1e-2)
+    solve!(qcp_smooth; max_iter=30, verbose=false, print_level=1)
+    
+    duration_before = sum(get_timesteps(get_trajectory(qcp_smooth)))
+    
+    # Convert to minimum-time problem
+    qcp_mintime = MinimumTimeProblem(qcp_smooth; final_fidelity=0.90, D=50.0)
+    
+    @test qcp_mintime isa QuantumControlProblem{<:EnsembleTrajectory}
+    @test qcp_mintime.qtraj isa EnsembleTrajectory{KetTrajectory}
+    
+    # Should have fidelity constraints for each ensemble member
+    # (one per trajectory in the ensemble)
+    
+    # Solve minimum-time problem
+    solve!(qcp_mintime; max_iter=30, verbose=false, print_level=1)
+    
+    duration_after = sum(get_timesteps(get_trajectory(qcp_mintime)))
+    
+    # Duration should not increase significantly
+    @test duration_after <= duration_before * 1.1
+
+    # Verify fidelity constraints are met for both states
+    traj = get_trajectory(qcp_mintime)
+    state_names = get_ensemble_state_names(qcp_mintime.qtraj)
+    goals = get_goal(qcp_mintime.qtraj)
+    
+    for (name, goal) in zip(state_names, goals)
+        ψ̃_final = traj[end][name]
+        ψ_final = iso_to_ket(ψ̃_final)
+        fid = fidelity(ψ_final, goal)
+        @test fid >= 0.89  # Just under constraint to account for numerical tolerance
+    end
+end
+
+@testitem "MinimumTimeProblem with EnsembleTrajectory (Unitary)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+    using LinearAlgebra
+
+    # Multi-goal minimum-time optimization
+    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    
+    # Create ensemble of unitary trajectories
+    # Note: Using same goal for both since different goals are very hard to optimize together
+    qtraj1 = UnitaryTrajectory(sys, GATES[:X], 30; Δt_bounds=(0.01, 0.5))
+    qtraj2 = UnitaryTrajectory(sys, GATES[:X], 30; Δt_bounds=(0.01, 0.5))
+    
+    ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+    
+    # Create and solve smooth pulse problem
+    qcp_smooth = SmoothPulseProblem(ensemble_qtraj; Q=100.0, R=1e-2)
+    solve!(qcp_smooth; max_iter=30, verbose=false, print_level=1)
+    
+    duration_before = sum(get_timesteps(get_trajectory(qcp_smooth)))
+    
+    # Convert to minimum-time problem
+    qcp_mintime = MinimumTimeProblem(qcp_smooth; final_fidelity=0.90, D=50.0)
+    
+    @test qcp_mintime isa QuantumControlProblem{<:EnsembleTrajectory}
+    @test qcp_mintime.qtraj isa EnsembleTrajectory{UnitaryTrajectory}
+    
+    # Solve minimum-time problem
+    solve!(qcp_mintime; max_iter=30, verbose=false, print_level=1)
+    
+    duration_after = sum(get_timesteps(get_trajectory(qcp_mintime)))
+    
+    # Duration should not increase significantly
+    @test duration_after <= duration_before * 1.1
+
+    # Verify fidelity constraints are met for all unitaries
+    traj = get_trajectory(qcp_mintime)
+    state_names = get_ensemble_state_names(qcp_mintime.qtraj)
+    goals = get_goal(qcp_mintime.qtraj)
+    
+    for (name, goal) in zip(state_names, goals)
+        Ũ⃗_final = traj[end][name]
+        U_final = iso_vec_to_operator(Ũ⃗_final)
+        fid = unitary_fidelity(U_final, goal)
+        @test fid >= 0.89  # Just under constraint to account for numerical tolerance
+    end
+end
+
+@testitem "MinimumTimeProblem with time-dependent UnitaryTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+    using LinearAlgebra
+
+    # Time-dependent Hamiltonian
+    ω = 2π * 5.0
+    H(u, t) = GATES[:Z] + u[1] * cos(ω * t) * GATES[:X] + u[2] * sin(ω * t) * GATES[:Y]
+    
+    sys = QuantumSystem(H, 1.0, [1.0, 1.0]; time_dependent=true)
+    
+    qtraj = UnitaryTrajectory(sys, GATES[:H], 30; Δt_bounds=(0.01, 0.5))
+    
+    # Verify time is in trajectory
+    @test haskey(qtraj.components, :t)
+    
+    # Create and solve smooth pulse problem
+    qcp_smooth = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+    
+    # Should have TimeIntegrator
+    has_time_integrator = any(int -> int isa TimeIntegrator, qcp_smooth.prob.integrators)
+    @test has_time_integrator
+    
+    solve!(qcp_smooth; max_iter=30, verbose=false, print_level=1)
+    
+    duration_before = sum(get_timesteps(get_trajectory(qcp_smooth)))
+    
+    # Convert to minimum-time
+    qcp_mintime = MinimumTimeProblem(qcp_smooth; final_fidelity=0.85, D=50.0)
+    
+    @test qcp_mintime isa QuantumControlProblem{UnitaryTrajectory}
+    
+    # Solve minimum-time problem
+    solve!(qcp_mintime; max_iter=30, verbose=false, print_level=1)
+    
+    duration_after = sum(get_timesteps(get_trajectory(qcp_mintime)))
+    @test duration_after <= duration_before * 1.1
+end
+
+@testitem "MinimumTimeProblem with time-dependent KetTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+    using LinearAlgebra
+
+    # Time-dependent Hamiltonian
+    ω = 2π * 5.0
+    H(u, t) = GATES[:Z] + u[1] * cos(ω * t) * GATES[:X]
+    
+    sys = QuantumSystem(H, 1.0, [1.0]; time_dependent=true)
+    
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+    
+    qtraj = KetTrajectory(sys, ψ_init, ψ_goal, 30; Δt_bounds=(0.01, 0.5))
+    
+    # Verify time is in trajectory
+    @test haskey(qtraj.components, :t)
+    
+    # Create and solve smooth pulse problem
+    qcp_smooth = SmoothPulseProblem(qtraj; Q=50.0, R=1e-3)
+    
+    # Should have TimeIntegrator
+    has_time_integrator = any(int -> int isa TimeIntegrator, qcp_smooth.prob.integrators)
+    @test has_time_integrator
+    
+    solve!(qcp_smooth; max_iter=30, verbose=false, print_level=1)
+    
+    duration_before = sum(get_timesteps(get_trajectory(qcp_smooth)))
+    
+    # Convert to minimum-time
+    qcp_mintime = MinimumTimeProblem(qcp_smooth; final_fidelity=0.85, D=50.0)
+    
+    @test qcp_mintime isa QuantumControlProblem{KetTrajectory}
+    
+    # Solve minimum-time problem
+    solve!(qcp_mintime; max_iter=30, verbose=false, print_level=1)
+    
+    duration_after = sum(get_timesteps(get_trajectory(qcp_mintime)))
+    @test duration_after <= duration_before * 1.1
+end
\ No newline at end of file
diff --git a/src/problem_templates/quantum_state_minimum_time_problem.jl b/src/problem_templates/quantum_state_minimum_time_problem.jl
deleted file mode 100644
index 68cd9aee..00000000
--- a/src/problem_templates/quantum_state_minimum_time_problem.jl
+++ /dev/null
@@ -1,114 +0,0 @@
-export QuantumStateMinimumTimeProblem
-
-
-"""
-    QuantumStateMinimumTimeProblem(traj, sys, obj, integrators, constraints; kwargs...)
-    QuantumStateMinimumTimeProblem(prob; kwargs...)
-
-Construct a `DirectTrajOptProblem` for the minimum time problem of reaching a target state.
-
-# Keyword Arguments
-- `state_name::Symbol=:ψ̃`: The symbol for the state variables.
-- `final_fidelity::Union{Real, Nothing}=nothing`: The final fidelity.
-- `D=1.0`: The cost weight on the time.
-- `piccolo_options::PiccoloOptions=PiccoloOptions()`: The Piccolo options.
-"""
-function QuantumStateMinimumTimeProblem end
-
-function QuantumStateMinimumTimeProblem(
-    trajectory::NamedTrajectory,
-    ψ_goals::AbstractVector{<:AbstractVector{<:ComplexF64}},
-    objective::Objective,
-    dynamics::TrajectoryDynamics,
-    constraints::Vector{<:AbstractConstraint};
-    state_name::Symbol=:ψ̃,
-    final_fidelity::Float64=1.0,
-    D=100.0,
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-)
-    state_names = [name for name in trajectory.names if startswith(string(name), string(state_name))]
-    @assert length(state_names) ≥ 1 "No matching states found in trajectory"
-    @assert length(state_names) == length(ψ_goals) "Number of states and goals must match"
-
-    if piccolo_options.verbose
-        println("    constructing QuantumStateMinimumTimeProblem...")
-        println("\tfinal fidelity: $(final_fidelity)")
-    end
-
-    objective += MinimumTimeObjective(trajectory, D=D)
-
-    for (state_name, ψ_goal) in zip(state_names, ψ_goals)
-        fidelity_constraint = FinalKetFidelityConstraint(
-            ψ_goal,
-            state_name,
-            final_fidelity,
-            trajectory,
-        )
-
-        push!(constraints, fidelity_constraint)
-    end
-
-    return DirectTrajOptProblem(
-        trajectory,
-        objective,
-        dynamics,
-        constraints
-    )
-end
-
-# TODO: goals from trajectory?
-
-function QuantumStateMinimumTimeProblem(
-    prob::DirectTrajOptProblem,
-    ψ_goals::AbstractVector{<:AbstractVector{<:ComplexF64}};
-    objective::Objective=prob.objective,
-    constraints::AbstractVector{<:AbstractConstraint}=prob.constraints,
-    kwargs...
-)
-    return QuantumStateMinimumTimeProblem(
-        prob.trajectory,
-        ψ_goals,
-        objective,
-        prob.dynamics,
-        constraints;
-        kwargs...
-    )
-end
-
-function QuantumStateMinimumTimeProblem(
-    prob::DirectTrajOptProblem,
-    ψ_goal::AbstractVector{<:ComplexF64};
-    kwargs...
-)
-    return QuantumStateMinimumTimeProblem(prob, [ψ_goal]; kwargs...)
-end
-
-# *************************************************************************** #
-
-@testitem "Test quantum state minimum time" begin
-    using NamedTrajectories
-    using PiccoloQuantumObjects 
-
-    T = 51
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    Δt = 0.2
-    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-    ψ_init = Vector{ComplexF64}[[1.0, 0.0]]
-    ψ_target = Vector{ComplexF64}[[0.0, 1.0]]
-
-    prob = QuantumStateSmoothPulseProblem(
-        sys, ψ_init, ψ_target, T, Δt;
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    initial = sum(get_timesteps(prob.trajectory))
-
-    min_prob = QuantumStateMinimumTimeProblem(
-        prob, ψ_target,
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    solve!(min_prob, max_iter=50, verbose=false, print_level=1)
-    final = sum(get_timesteps(min_prob.trajectory))
-
-    @test final < initial
-end
diff --git a/src/problem_templates/quantum_state_sampling_problem.jl b/src/problem_templates/quantum_state_sampling_problem.jl
deleted file mode 100644
index bd07d7f4..00000000
--- a/src/problem_templates/quantum_state_sampling_problem.jl
+++ /dev/null
@@ -1,237 +0,0 @@
-export QuantumStateSamplingProblem
-
-"""
-
-"""
-function QuantumStateSamplingProblem end
-
-function QuantumStateSamplingProblem(
-    systems::AbstractVector{<:AbstractQuantumSystem},
-    ψ_inits::AbstractVector{<:AbstractVector{<:AbstractVector{<:ComplexF64}}},
-    ψ_goals::AbstractVector{<:AbstractVector{<:AbstractVector{<:ComplexF64}}},
-    T::Int,
-    Δt::Union{Float64, AbstractVector{Float64}};
-    ket_integrator=KetIntegrator,
-    system_weights=fill(1.0, length(systems)),
-    init_trajectory::Union{NamedTrajectory,Nothing}=nothing,
-    state_name::Symbol=:ψ̃,
-    control_name::Symbol = :u,
-    timestep_name::Symbol=:Δt,
-    u_guess::Union{Matrix{Float64},Nothing}=nothing,
-    du_bound::Float64 = Inf,
-    du_bounds = fill(du_bound, systems[1].n_drives),
-    ddu_bound::Float64 = 1.0,
-    ddu_bounds = fill(ddu_bound, systems[1].n_drives),
-    Δt_min::Float64=0.5 * minimum(Δt),
-    Δt_max::Float64=2.0 * maximum(Δt),
-    Q::Float64=100.0,
-    R=1e-2,
-    R_u::Union{Float64,Vector{Float64}} = R,
-    R_du::Union{Float64,Vector{Float64}} = R,
-    R_ddu::Union{Float64,Vector{Float64}} = R,
-    state_leakage_indices::Union{Nothing,AbstractVector{Int}}=nothing,
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-)
-    @assert length(ψ_inits) == length(ψ_goals)
-
-    if piccolo_options.verbose
-        println("    constructing QuantumStateSamplingProblem...")
-        println("\tusing integrator: $(typeof(ket_integrator))")
-        println("\tusing $(length(ψ_inits)) initial state(s)")
-    end
-
-    # Outer dimension is the system, inner dimension is the initial state
-    state_names = [
-        [Symbol(string(state_name) * "$(i)_system_$(j)") for i in eachindex(ψs)]
-        for (j, ψs) ∈ zip(eachindex(systems), ψ_inits)
-    ]
-        
-    # Trajectory
-    if !isnothing(init_trajectory)
-        traj = init_trajectory
-    else
-        # TODO(main issue with this is that it doesn't handle composite systems correctly, will come back to this.
-        trajs = map(zip(systems, state_names, ψ_inits, ψ_goals)) do (sys, names, ψis, ψgs)
-            initialize_trajectory(
-                ψis,
-                ψgs,
-                T,
-                Δt,
-                sys.n_drives,
-                (sys.drive_bounds, du_bounds, ddu_bounds);
-                state_names=names,
-                control_name=control_name,
-                timestep_name=timestep_name,
-                Δt_bounds=(Δt_min, Δt_max),
-                zero_initial_and_final_derivative=piccolo_options.zero_initial_and_final_derivative,
-                bound_state=piccolo_options.bound_state,
-                u_guess = u_guess,
-                system=sys,
-                rollout_integrator=piccolo_options.rollout_integrator,
-                verbose=false, # loop
-            )
-        end
-
-        traj = merge(
-            trajs,
-            merge_names = (u = 1, du = 1, ddu = 1, Δt = 1),
-            timestep = timestep_name,
-        )
-    end
-
-    control_names = [
-        name for name ∈ traj.names
-        if endswith(string(name), string(control_name))
-    ]
-
-    # Objective
-    J = QuadraticRegularizer(control_names[1], traj, R_u)
-    J += QuadraticRegularizer(control_names[2], traj, R_du)
-    J += QuadraticRegularizer(control_names[3], traj, R_ddu)
-
-    for (weight, names) in zip(system_weights, state_names)
-        for name in names
-            J += KetInfidelityObjective(name, traj, Q=weight * Q)
-        end
-    end
-
-    # Optional Piccolo constraints and objectives
-    J += apply_piccolo_options!(
-        piccolo_options, constraints, traj;
-        state_names=vcat(state_names...),
-        state_leakage_indices=isnothing(state_leakage_indices) ? nothing : fill(state_leakage_indices, length(vcat(state_names...))),
-    )
-
-    # Integrators
-    state_integrators = []
-    for (sys, names) ∈ zip(systems, state_names)
-        for name ∈ names
-            push!(state_integrators, ket_integrator(sys, traj, name, control_name))
-        end
-    end
-
-    integrators = [
-        state_integrators...,
-        DerivativeIntegrator(traj, control_name, control_names[2]),
-        DerivativeIntegrator(traj, control_names[2], control_names[3]),
-    ]
-
-    return DirectTrajOptProblem(
-        traj,
-        J,
-        integrators;
-        constraints=constraints,
-    )
-end
-
-function QuantumStateSamplingProblem(
-    systems::AbstractVector{<:AbstractQuantumSystem},
-    ψ_inits::AbstractVector{<:AbstractVector{<:ComplexF64}},
-    ψ_goals::AbstractVector{<:AbstractVector{<:ComplexF64}},
-    args...;
-    kwargs...
-)
-    return QuantumStateSamplingProblem(
-        systems, 
-        fill(ψ_inits, length(systems)), 
-        fill(ψ_goals, length(systems)),
-        args...; 
-        kwargs...
-    )
-end
-
-function QuantumStateSamplingProblem(
-    systems::AbstractVector{<:AbstractQuantumSystem},
-    ψ_init::AbstractVector{<:ComplexF64},
-    ψ_goal::AbstractVector{<:ComplexF64},
-    args...;
-    kwargs...
-)
-    # Broadcast the initial and target states to all systems
-    return QuantumStateSamplingProblem(systems, [ψ_init], [ψ_goal], args...; kwargs...)
-end
-
-# *************************************************************************** #
-
-@testitem "Sample systems with single initial, target" begin
-    using PiccoloQuantumObjects
-
-    T = 50
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    Δt = 0.2
-    sys1 = QuantumSystem(0.3 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-    sys2 = QuantumSystem(-0.3 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-    ψ_init = Vector{ComplexF64}([1.0, 0.0])
-    ψ_target = Vector{ComplexF64}([0.0, 1.0])
-    
-    prob = QuantumStateSamplingProblem(
-        [sys1, sys2], ψ_init, ψ_target, T, Δt;
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    
-    state_name = :ψ̃
-    state_names = [n for n ∈ prob.trajectory.names if startswith(string(n), string(state_name))]
-    sys_state_names = [n for n ∈ state_names if endswith(string(n), "1")]
-    
-    # Separately compute all unique initial and goal state fidelities
-    inits = []
-    for sys in [sys1, sys2]
-        push!(
-            inits,
-            [rollout_fidelity(prob.trajectory, sys; state_name = n) for n in state_names],
-        )
-    end
-    
-    solve!(prob, max_iter=20, print_level=1, verbose=false)
-        
-    for (init, sys) in zip(inits, [sys1, sys2])
-        final =
-            [rollout_fidelity(prob.trajectory, sys, state_name = n) for n in state_names]
-        @test all(final .> init)
-    end
-end
-
-@testitem "Sample systems with multiple initial, target" begin
-    using PiccoloQuantumObjects
-
-    T = 50
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    Δt = 0.2
-    sys1 = QuantumSystem(0.3 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-    sys2 = QuantumSystem(-0.3 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-
-    # Multiple initial and target states
-    ψ_inits = Vector{ComplexF64}.([[1.0, 0.0], [0.0, 1.0]])
-    ψ_targets = Vector{ComplexF64}.([[0.0, 1.0], [1.0, 0.0]])
-
-    prob = QuantumStateSamplingProblem(
-        [sys1, sys2], ψ_inits, ψ_targets, T, Δt;
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-
-    state_name = :ψ̃
-    state_names = [n for n ∈ prob.trajectory.names if startswith(string(n), string(state_name))]
-    sys_state_names = [n for n ∈ state_names if endswith(string(n), "1")]
-
-    # Separately compute all unique initial and goal state fidelities
-    inits = []
-    for sys in [sys1, sys2]
-        push!(
-            inits,
-            [rollout_fidelity(prob.trajectory, sys; state_name = n) for n in state_names],
-        )
-    end
-
-    solve!(prob, max_iter=20, print_level=1, verbose=false)
-
-    for (init, sys) in zip(inits, [sys1, sys2])
-        final =
-            [rollout_fidelity(prob.trajectory, sys, state_name = n) for n in state_names]
-        @test all(final .> init)
-    end
-end
-
-# TODO: Test that u_guess can be used
diff --git a/src/problem_templates/quantum_state_smooth_pulse_problem.jl b/src/problem_templates/quantum_state_smooth_pulse_problem.jl
deleted file mode 100644
index 8d19e09e..00000000
--- a/src/problem_templates/quantum_state_smooth_pulse_problem.jl
+++ /dev/null
@@ -1,232 +0,0 @@
-export QuantumStateSmoothPulseProblem
-
-"""
-    QuantumStateSmoothPulseProblem(system, ψ_inits, ψ_goals, T, Δt; kwargs...)
-    QuantumStateSmoothPulseProblem(system, ψ_init, ψ_goal, T, Δt; kwargs...)
-    QuantumStateSmoothPulseProblem(H_drift, H_drives, args...; kwargs...)
-
-Create a quantum state smooth pulse problem. The goal is to find a control pulse `u(t)`
-that drives all of the initial states `ψ_inits` to the corresponding target states
-`ψ_goals` using `T` timesteps of size `Δt`. This problem also controls the  first and
-second derivatives of the control pulse, `du(t)` and `ddu(t)`, to ensure smoothness.
-
-# Arguments
-- `system::AbstractQuantumSystem`: The quantum system.
-or
-- `H_drift::AbstractMatrix{<:Number}`: The drift Hamiltonian.
-- `H_drives::Vector{<:AbstractMatrix{<:Number}}`: The control Hamiltonians.
-with
-- `ψ_inits::Vector{<:AbstractVector{<:ComplexF64}}`: The initial states.
-- `ψ_goals::Vector{<:AbstractVector{<:ComplexF64}}`: The target states.
-or
-- `ψ_init::AbstractVector{<:ComplexF64}`: The initial state.
-- `ψ_goal::AbstractVector{<:ComplexF64}`: The target state.
-with
-- `T::Int`: The number of timesteps.
-- `Δt::Float64`: The timestep size.
-
-
-# Keyword Arguments
-- `state_name::Symbol=:ψ̃`: The name of the state variable.
-- `control_name::Symbol=:u`: The name of the control variable.
-- `timestep_name::Symbol=:Δt`: The name of the timestep variable.
-- `init_trajectory::Union{NamedTrajectory, Nothing}=nothing`: The initial trajectory.
-- `u_bound::Float64=1.0`: The bound on the control pulse.
-- `u_bounds=fill(u_bound, length(system.G_drives))`: The bounds on the control pulse.
-- `u_guess::Union{Matrix{Float64}, Nothing}=nothing`: The initial guess for the control pulse.
-- `du_bound::Float64=Inf`: The bound on the first derivative of the control pulse.
-- `du_bounds=fill(du_bound, length(system.G_drives))`: The bounds on the first derivative of the control pulse.
-- `ddu_bound::Float64=1.0`: The bound on the second derivative of the control pulse.
-- `ddu_bounds=fill(ddu_bound, length(system.G_drives))`: The bounds on the second derivative of the control pulse.
-- `Δt_min::Float64=0.5 * Δt`: The minimum timestep size.
-- `Δt_max::Float64=1.5 * Δt`: The maximum timestep size.
-- `drive_derivative_σ::Float64=0.01`: The standard deviation of the drive derivative random initialization.
-- `Q::Float64=100.0`: The weight on the state objective.
-- `R=1e-2`: The weight on the control pulse and its derivatives.
-- `R_u::Union{Float64, Vector{Float64}}=R`: The weight on the control pulse.
-- `R_du::Union{Float64, Vector{Float64}}=R`: The weight on the first derivative of the control pulse.
-- `R_ddu::Union{Float64, Vector{Float64}}=R`: The weight on the second derivative of the control pulse.
-- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: The constraints.
-- `piccolo_options::PiccoloOptions=PiccoloOptions()`: The Piccolo options.
-"""
-function QuantumStateSmoothPulseProblem end
-
-function QuantumStateSmoothPulseProblem(
-    sys::AbstractQuantumSystem,
-    ψ_inits::Vector{<:AbstractVector{<:ComplexF64}},
-    ψ_goals::Vector{<:AbstractVector{<:ComplexF64}},
-    T::Int,
-    Δt::Union{Float64, <:AbstractVector{Float64}};
-    ket_integrator=KetIntegrator,
-    state_name::Symbol=:ψ̃,
-    control_name::Symbol=:u,
-    timestep_name::Symbol=:Δt,
-    init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
-    u_guess::Union{AbstractMatrix{Float64}, Nothing}=nothing,
-    du_bound::Float64=Inf,
-    du_bounds=fill(du_bound, sys.n_drives),
-    ddu_bound::Float64=1.0,
-    ddu_bounds=fill(ddu_bound, sys.n_drives),
-    Δt_min::Float64=0.5 * minimum(Δt),
-    Δt_max::Float64=2.0 * maximum(Δt),
-    Q::Float64=100.0,
-    R=1e-2,
-    R_u::Union{Float64, Vector{Float64}}=R,
-    R_du::Union{Float64, Vector{Float64}}=R,
-    R_ddu::Union{Float64, Vector{Float64}}=R,
-    state_leakage_indices::Union{Nothing, AbstractVector{Int}}=nothing,
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-)
-    @assert length(ψ_inits) == length(ψ_goals)
-
-    if piccolo_options.verbose
-        println("    constructing QuantumStateSmoothPulseProblem...")
-        println("\tusing integrator: $(typeof(ket_integrator))")
-        println("\tusing $(length(ψ_inits)) initial state(s)")
-    end
-
-    # Trajectory
-    if !isnothing(init_trajectory)
-        traj = init_trajectory
-    else
-        traj = initialize_trajectory(
-            ψ_goals,
-            ψ_inits,
-            T,
-            Δt,
-            sys.n_drives,
-            (sys.drive_bounds, du_bounds, ddu_bounds);
-            state_name=state_name,
-            control_name=control_name,
-            timestep_name=timestep_name,
-            zero_initial_and_final_derivative=piccolo_options.zero_initial_and_final_derivative,
-            Δt_bounds=(Δt_min, Δt_max),
-            bound_state=piccolo_options.bound_state,
-            u_guess=u_guess,
-            system=sys,
-            rollout_integrator=piccolo_options.rollout_integrator,
-        )
-    end
-
-    state_names = [
-        name for name ∈ traj.names
-            if startswith(string(name), string(state_name))
-    ]
-    @assert length(state_names) == length(ψ_inits) "Number of states must match number of initial states"
-
-    control_names = [
-        name for name ∈ traj.names
-            if endswith(string(name), string(control_name))
-    ]
-
-    # Objective
-    J = QuadraticRegularizer(control_names[1], traj, R_u)
-    J += QuadraticRegularizer(control_names[2], traj, R_du)
-    J += QuadraticRegularizer(control_names[3], traj, R_ddu)
-
-    for name ∈ state_names
-        J += KetInfidelityObjective(name, traj; Q=Q)
-    end
-
-    # Optional Piccolo constraints and objectives
-    J += apply_piccolo_options!(
-        piccolo_options, constraints, traj;
-        state_names=state_names,
-        state_leakage_indices= isnothing(state_leakage_indices) ? nothing : fill(state_leakage_indices, length(state_names))
-    )
-
-    state_names = [
-        name for name ∈ traj.names
-            if startswith(string(name), string(state_name))
-    ]
-
-    state_integrators = []
-
-    for name ∈ state_names
-        push!(
-            state_integrators, 
-            ket_integrator(sys, traj, name, control_name)
-        )
-    end
-
-    integrators = [
-        state_integrators...,
-        DerivativeIntegrator(traj, control_name, control_names[2]),
-        DerivativeIntegrator(traj, control_names[2], control_names[3])
-    ]
-
-    return DirectTrajOptProblem(
-        traj,
-        J,
-        integrators;
-        constraints=constraints,
-    )
-end
-
-function QuantumStateSmoothPulseProblem(
-    system::AbstractQuantumSystem,
-    ψ_init::AbstractVector{<:ComplexF64},
-    ψ_goal::AbstractVector{<:ComplexF64},
-    args...;
-    kwargs...
-)
-    return QuantumStateSmoothPulseProblem(system, [ψ_init], [ψ_goal], args...; kwargs...)
-end
-
-function QuantumStateSmoothPulseProblem(
-    H_drift::AbstractMatrix{<:Number},
-    H_drives::Vector{<:AbstractMatrix{<:Number}},
-    args...;
-    kwargs...
-)
-    system = QuantumSystem(H_drift, H_drives, T_max, u_bounds)
-    return QuantumStateSmoothPulseProblem(system, args...; kwargs...)
-end
-
-# *************************************************************************** #
-
-@testitem "Test quantum state smooth pulse" begin
-    using PiccoloQuantumObjects 
-
-    T = 51
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    Δt = 0.2
-    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-    ψ_init = Vector{ComplexF64}([1.0, 0.0])
-    ψ_target = Vector{ComplexF64}([0.0, 1.0])
-    
-    # Single initial and target states
-    # --------------------------------
-    prob = QuantumStateSmoothPulseProblem(
-        sys, ψ_init, ψ_target, T, Δt;
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    initial = rollout_fidelity(prob.trajectory, sys, control_name=:u)
-    solve!(prob, max_iter=50, print_level=1, verbose=false)
-    final = rollout_fidelity(prob.trajectory, sys, control_name=:u)
-    @test final > initial
-end
-
-@testitem "Test multiple quantum states smooth pulse" begin
-    using PiccoloQuantumObjects 
-
-    T = 50
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    Δt = 0.2
-    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-    ψ_inits = Vector{ComplexF64}.([[1.0, 0.0], [0.0, 1.0]])
-    ψ_targets = Vector{ComplexF64}.([[0.0, 1.0], [1.0, 0.0]])
-
-    prob = QuantumStateSmoothPulseProblem(
-        sys, ψ_inits, ψ_targets, T, Δt;
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    initial = rollout_fidelity(prob.trajectory, sys, control_name=:u)
-    solve!(prob, max_iter=50, print_level=1, verbose=false)
-    final = rollout_fidelity(prob.trajectory, sys, control_name=:u)
-    final, initial
-    @test all(final .> initial)
-end
diff --git a/src/problem_templates/sampling_problem.jl b/src/problem_templates/sampling_problem.jl
new file mode 100644
index 00000000..93ead06d
--- /dev/null
+++ b/src/problem_templates/sampling_problem.jl
@@ -0,0 +1,371 @@
+export SamplingProblem
+
+# Note: SamplingTrajectory is now exported from PiccoloQuantumObjects
+
+# ============================================================================= #
+# Sampling Objective Utilities
+# ============================================================================= #
+
+"""
+    extract_regularization(objective, state_sym::Symbol) -> AbstractObjective
+
+Extract regularization terms (non-state-dependent objectives) from a composite objective.
+"""
+function extract_regularization(objective, state_sym::Symbol)
+    objs = hasproperty(objective, :objectives) ? objective.objectives : [objective]
+    
+    regularizers = filter(objs) do term
+        term_syms = if hasproperty(term, :syms)
+            term.syms
+        elseif hasproperty(term, :var_names)
+            term.var_names
+        else
+            Symbol[]
+        end
+        state_sym ∉ term_syms
+    end
+    
+    return isempty(regularizers) ? NullObjective() : reduce(+, regularizers)
+end
+
+# ============================================================================= #
+# Sampling State Objective (dispatch-based)
+# ============================================================================= #
+
+"""
+    sampling_state_objective(qtraj, traj, state_sym, Q)
+
+Create the state-dependent objective for a sampling member.
+Dispatches on quantum trajectory type.
+"""
+function sampling_state_objective(
+    qtraj::UnitaryTrajectory, 
+    traj::NamedTrajectory, 
+    state_sym::Symbol, 
+    Q::Float64
+)
+    return UnitaryInfidelityObjective(get_goal(qtraj), state_sym, traj; Q=Q)
+end
+
+function sampling_state_objective(
+    qtraj::KetTrajectory, 
+    traj::NamedTrajectory, 
+    state_sym::Symbol, 
+    Q::Float64
+)
+    ψ_goal = get_goal(qtraj)
+    return KetInfidelityObjective(ψ_goal, state_sym, traj; Q=Q)
+end
+
+function sampling_state_objective(
+    qtraj::DensityTrajectory, 
+    traj::NamedTrajectory, 
+    state_sym::Symbol, 
+    Q::Float64
+)
+    # DensityTrajectory doesn't have a fidelity objective yet
+    # Return NullObjective for now
+    return NullObjective(traj)
+end
+
+# ============================================================================= #
+# SamplingProblem Constructor
+# ============================================================================= #
+
+@doc raw"""
+    SamplingProblem(qcp::QuantumControlProblem, systems::Vector{<:AbstractQuantumSystem}; kwargs...)
+
+Construct a `SamplingProblem` from an existing `QuantumControlProblem` and a list of systems.
+
+This creates a robust optimization problem where the controls are shared across all systems,
+but each system evolves according to its own dynamics. The objective is the weighted sum of
+fidelity objectives for each system.
+
+# Arguments
+- `qcp::QuantumControlProblem`: The base problem (defines nominal trajectory, objective, etc.)
+- `systems::Vector{<:AbstractQuantumSystem}`: List of systems to optimize over
+
+# Keyword Arguments
+- `weights::Vector{Float64}=fill(1.0, length(systems))`: Weights for each system
+- `Q::Float64=100.0`: Weight on infidelity objective (explicit, not extracted from base problem)
+- `piccolo_options::PiccoloOptions=PiccoloOptions()`: Options for the solver
+
+# Returns
+- `QuantumControlProblem{SamplingTrajectory}`: A new problem with the sampling trajectory
+"""
+function SamplingProblem(
+    qcp::QuantumControlProblem,
+    systems::Vector{<:AbstractQuantumSystem};
+    weights::Vector{Float64}=fill(1.0, length(systems)),
+    Q::Float64=100.0,
+    piccolo_options::PiccoloOptions=PiccoloOptions(),
+)
+    if piccolo_options.verbose
+        println("    constructing SamplingProblem...")
+        println("\tusing $(length(systems)) systems")
+    end
+
+    base_qtraj = qcp.qtraj
+    state_sym = get_state_name(base_qtraj)
+
+    # 1. Build sampling trajectory with duplicated states
+    new_traj, state_names = build_sampling_trajectory(
+        get_trajectory(qcp), 
+        state_sym, 
+        length(systems)
+    )
+
+    # 2. Create SamplingTrajectory wrapper
+    sampling_qtraj = SamplingTrajectory(base_qtraj, systems, weights, state_names)
+
+    # 3. Build objective: weighted state objectives + shared regularization
+    J_state = sum(
+        sampling_state_objective(base_qtraj, new_traj, name, w * Q)
+        for (name, w) in zip(state_names, weights)
+    )
+    J_reg = extract_regularization(qcp.prob.objective, state_sym)
+    J_total = J_state + J_reg
+
+    # 4. Build integrators: shared (derivative, time) + dynamics for each system
+    shared_integrators = filter(qcp.prob.integrators) do int
+        int isa DerivativeIntegrator || int isa TimeIntegrator
+    end
+    
+    # Use BilinearIntegrator dispatch on SamplingTrajectory
+    dynamics_integrators = BilinearIntegrator(sampling_qtraj, new_traj)
+    
+    all_integrators = vcat(shared_integrators, dynamics_integrators)
+
+    # 5. Construct problem
+    prob = DirectTrajOptProblem(
+        new_traj,
+        J_total,
+        all_integrators;
+        constraints=AbstractConstraint[]
+    )
+
+    return QuantumControlProblem(sampling_qtraj, prob)
+end
+
+# ============================================================================= #
+# Composability with MinimumTimeProblem
+# ============================================================================= #
+
+function _update_goal(qtraj::SamplingTrajectory, new_goal)
+    new_base = _update_goal(qtraj.base_trajectory, new_goal)
+    return update_base_trajectory(qtraj, new_base)
+end
+
+function _final_fidelity_constraint(
+    qtraj::SamplingTrajectory,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    constraints = [
+        _sampling_fidelity_constraint(qtraj.base_trajectory, name, final_fidelity, traj)
+        for name in get_ensemble_state_names(qtraj)
+    ]
+    return constraints
+end
+
+# Dispatch on trajectory type for fidelity constraint
+function _sampling_fidelity_constraint(
+    qtraj::UnitaryTrajectory,
+    state_sym::Symbol,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    return FinalUnitaryFidelityConstraint(get_goal(qtraj), state_sym, final_fidelity, traj)
+end
+
+function _sampling_fidelity_constraint(
+    qtraj::KetTrajectory,
+    state_sym::Symbol,
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    return FinalKetFidelityConstraint(get_goal(qtraj), state_sym, final_fidelity, traj)
+end
+
+# Tests
+@testitem "SamplingProblem Construction" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Define system
+    sys = QuantumSystem(GATES[:Z], [GATES[:X]], 10.0, [1.0])
+
+    # Create base problem
+    qtraj = UnitaryTrajectory(sys, GATES[:H], 10)
+    qcp = SmoothPulseProblem(qtraj; Q=100.0)
+
+    # Create sampling problem with 2 systems
+    systems = [sys, sys] # Identical systems for testing
+    sampling_prob = SamplingProblem(qcp, systems)
+
+    @test sampling_prob isa QuantumControlProblem
+    @test sampling_prob.qtraj isa SamplingTrajectory
+    @test length(sampling_prob.qtraj.systems) == 2
+
+    # Check trajectory components (now use _sample_ suffix)
+    traj = get_trajectory(sampling_prob)
+    @test haskey(traj.components, :Ũ⃗_sample_1)
+    @test haskey(traj.components, :Ũ⃗_sample_2)
+    @test haskey(traj.components, :u)
+
+    # Check integrators
+    # Should have 2 dynamics integrators + derivative integrators
+    # SmoothPulseProblem adds 2 derivative integrators.
+    # So total should be 2 + 2 = 4 (plus maybe time integrator if time dependent)
+    @test length(sampling_prob.prob.integrators) >= 4
+end
+
+@testitem "SamplingProblem Solving" tags = [:sampling_problem] begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Simple robust optimization
+    # System with uncertainty in drift
+    sys_nominal = QuantumSystem(GATES[:Z], [GATES[:X]], 10.0, [1.0])
+    sys_perturbed = QuantumSystem(1.1 * GATES[:Z], [GATES[:X]], 10.0, [1.0])
+
+    qtraj = UnitaryTrajectory(sys_nominal, GATES[:X], 20)
+    qcp = SmoothPulseProblem(qtraj; Q=100.0)
+
+    sampling_prob = SamplingProblem(qcp, [sys_nominal, sys_perturbed])
+
+    # Solve
+    solve!(sampling_prob; max_iter=10, verbose=false)
+
+    # Check that we have a solution
+    @test sampling_prob.prob.objective(sampling_prob.trajectory) < 1e10 # Just check it didn't blow up
+end
+
+@testitem "SamplingProblem with KetTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Robust state transfer over parameter uncertainty
+    sys_nominal = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    sys_perturbed = QuantumSystem(1.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+    qtraj = KetTrajectory(sys_nominal, ψ_init, ψ_goal, 20)
+
+    qcp = SmoothPulseProblem(qtraj; Q=50.0, R=1e-3)
+
+    # Create sampling problem
+    sampling_prob = SamplingProblem(qcp, [sys_nominal, sys_perturbed]; Q=50.0)
+
+    @test sampling_prob isa QuantumControlProblem
+    @test sampling_prob.qtraj isa SamplingTrajectory
+
+    # Check trajectory has sample states
+    traj = get_trajectory(sampling_prob)
+    @test haskey(traj.components, :ψ̃_sample_1)
+    @test haskey(traj.components, :ψ̃_sample_2)
+
+    # Solve
+    solve!(sampling_prob; max_iter=10, verbose=false, print_level=1)
+end
+
+@testitem "SamplingProblem with custom weights" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    sys1 = QuantumSystem(GATES[:Z], [GATES[:X]], 10.0, [1.0])
+    sys2 = QuantumSystem(1.1 * GATES[:Z], [GATES[:X]], 10.0, [1.0])
+    sys3 = QuantumSystem(0.9 * GATES[:Z], [GATES[:X]], 10.0, [1.0])
+
+    qtraj = UnitaryTrajectory(sys1, GATES[:X], 15)
+    qcp = SmoothPulseProblem(qtraj; Q=100.0)
+
+    # Non-uniform weights - emphasize nominal system
+    weights = [0.6, 0.2, 0.2]
+    sampling_prob = SamplingProblem(qcp, [sys1, sys2, sys3]; weights=weights, Q=100.0)
+
+    @test sampling_prob.qtraj.weights == weights
+    @test length(sampling_prob.qtraj.systems) == 3
+
+    # Should have 3 sample states
+    traj = get_trajectory(sampling_prob)
+    @test haskey(traj.components, :Ũ⃗_sample_1)
+    @test haskey(traj.components, :Ũ⃗_sample_2)
+    @test haskey(traj.components, :Ũ⃗_sample_3)
+
+    solve!(sampling_prob; max_iter=5, verbose=false, print_level=1)
+end
+
+@testitem "SamplingProblem + MinimumTimeProblem composition" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Robust minimum-time optimization
+    sys_nominal = QuantumSystem(0.1 * GATES[:Z], [GATES[:X]], 1.0, [1.0])
+    sys_perturbed = QuantumSystem(0.11 * GATES[:Z], [GATES[:X]], 1.0, [1.0])
+
+    qtraj = UnitaryTrajectory(sys_nominal, GATES[:X], 30; Δt_bounds=(0.01, 0.5))
+    qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+
+    # First create sampling problem
+    sampling_prob = SamplingProblem(qcp, [sys_nominal, sys_perturbed]; Q=100.0)
+    solve!(sampling_prob; max_iter=20, verbose=false, print_level=1)
+
+    # Then convert to minimum-time
+    mintime_prob = MinimumTimeProblem(sampling_prob; final_fidelity=0.90, D=50.0)
+
+    @test mintime_prob isa QuantumControlProblem
+    @test mintime_prob.qtraj isa SamplingTrajectory
+
+    # Solve minimum-time
+    solve!(mintime_prob; max_iter=20, verbose=false, print_level=1)
+end
+@testitem "SamplingProblem with DensityTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    # Robust open-system control over parameter uncertainty
+    sys_nominal = OpenQuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    sys_perturbed = OpenQuantumSystem(1.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+
+    ρ_init = ComplexF64[1.0 0.0; 0.0 0.0]  # |0⟩⟨0|
+    ρ_goal = ComplexF64[0.0 0.0; 0.0 1.0]  # |1⟩⟨1|
+    
+    qtraj = DensityTrajectory(sys_nominal, ρ_init, ρ_goal, 20)
+
+    qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-3)
+
+    # Create sampling problem
+    sampling_prob = SamplingProblem(qcp, [sys_nominal, sys_perturbed]; Q=100.0)
+
+    @test sampling_prob isa QuantumControlProblem
+    @test sampling_prob.qtraj isa SamplingTrajectory{DensityTrajectory}
+
+    # Check trajectory has sample states
+    traj = get_trajectory(sampling_prob)
+    @test haskey(traj.components, :ρ⃗̃_sample_1)
+    @test haskey(traj.components, :ρ⃗̃_sample_2)
+
+    # Check integrators (2 dynamics + 2 derivatives)
+    @test length(sampling_prob.prob.integrators) >= 4
+
+    # Solve and verify dynamics are satisfied
+    solve!(sampling_prob; max_iter=20, verbose=false, print_level=1)
+    
+    # Test dynamics constraints are satisfied for all integrators
+    for integrator in sampling_prob.prob.integrators
+        if integrator isa BilinearIntegrator
+            δ = zeros(integrator.dim)
+            DirectTrajOpt.evaluate!(δ, integrator, traj)
+            @test norm(δ, Inf) < 1e-3
+        end
+    end
+end
\ No newline at end of file
diff --git a/src/problem_templates/smooth_pulse_problem.jl b/src/problem_templates/smooth_pulse_problem.jl
new file mode 100644
index 00000000..5ca1ce48
--- /dev/null
+++ b/src/problem_templates/smooth_pulse_problem.jl
@@ -0,0 +1,987 @@
+export SmoothPulseProblem
+
+@doc raw"""
+    SmoothPulseProblem(qtraj::AbstractQuantumTrajectory; kwargs...)
+
+Construct a `QuantumControlProblem` for smooth pulse optimization that dispatches on the quantum trajectory type.
+
+This unified interface replaces `UnitarySmoothPulseProblem`, `QuantumStateSmoothPulseProblem`, etc.
+The problem automatically:
+- Adds control derivatives (u̇, ü) to the trajectory
+- Creates appropriate dynamics integrator via `default_integrator`
+- Adds derivative integrators for smoothness constraints
+- Constructs objective with infidelity and regularization terms
+
+# Arguments
+- `qtraj::AbstractQuantumTrajectory`: Quantum trajectory (UnitaryTrajectory, KetTrajectory, or DensityTrajectory)
+
+# Keyword Arguments
+- `integrator::Union{Nothing, AbstractIntegrator, Vector{<:AbstractIntegrator}}=nothing`: Optional custom integrator(s). If `nothing`, uses default type-appropriate integrator. For KetTrajectory with multiple states, this will automatically return a vector of integrators (one per state). Can also be explicitly provided as a vector.
+- `du_bound::Float64=Inf`: Bound on first derivative
+- `ddu_bound::Float64=1.0`: Bound on second derivative
+- `Q::Float64=100.0`: Weight on infidelity/objective
+- `R::Float64=1e-2`: Weight on regularization terms (u, u̇, ü)
+- `R_u::Union{Float64, Vector{Float64}}=R`: Weight on control regularization
+- `R_du::Union{Float64, Vector{Float64}}=R`: Weight on first derivative regularization
+- `R_ddu::Union{Float64, Vector{Float64}}=R`: Weight on second derivative regularization
+- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: Additional constraints
+- `piccolo_options::PiccoloOptions=PiccoloOptions()`: Piccolo solver options
+
+# Returns
+- `QuantumControlProblem`: Wrapper containing quantum trajectory and optimization problem
+
+# Examples
+```julia
+# Unitary gate synthesis
+sys = QuantumSystem(H_drift, H_drives, T, drive_bounds)
+qtraj = UnitaryTrajectory(sys, U_goal, N)
+qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+solve!(qcp; max_iter=100)
+
+# Quantum state transfer
+qtraj = KetTrajectory(sys, ψ_init, ψ_goal, N)
+qcp = SmoothPulseProblem(qtraj; Q=50.0, R=1e-3)
+solve!(qcp)
+
+# Open system
+open_sys = OpenQuantumSystem(H_drift, H_drives, T, drive_bounds)
+qtraj = DensityTrajectory(open_sys, ρ_init, ρ_goal, N)
+qcp = SmoothPulseProblem(qtraj; Q=100.0)
+solve!(qcp)
+```
+"""
+function SmoothPulseProblem(
+    qtraj::AbstractQuantumTrajectory;
+    integrator::Union{Nothing,AbstractIntegrator,Vector{<:AbstractIntegrator}}=nothing,
+    du_bound::Float64=Inf,
+    ddu_bound::Float64=1.0,
+    Q::Float64=100.0,
+    R::Float64=1e-2,
+    R_u::Union{Float64,Vector{Float64}}=R,
+    R_du::Union{Float64,Vector{Float64}}=R,
+    R_ddu::Union{Float64,Vector{Float64}}=R,
+    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
+    piccolo_options::PiccoloOptions=PiccoloOptions(),
+)
+    if piccolo_options.verbose
+        println("    constructing SmoothPulseProblem for $(typeof(qtraj))...")
+    end
+
+    # Extract info from quantum trajectory
+    sys = PiccoloQuantumObjects.get_system(qtraj)
+    state_sym = PiccoloQuantumObjects.get_state_name(qtraj)
+    control_sym = PiccoloQuantumObjects.get_control_name(qtraj)
+
+    # Add control derivatives to trajectory (always 2 derivatives for smooth pulses)
+    du_bounds = fill(du_bound, sys.n_drives)
+    ddu_bounds = fill(ddu_bound, sys.n_drives)
+
+    traj_smooth = add_control_derivatives(
+        PiccoloQuantumObjects.get_trajectory(qtraj),
+        2;  # Always use 2 derivatives
+        control_name=control_sym,
+        derivative_bounds=(du_bounds, ddu_bounds)
+    )
+
+    # Get control derivative names
+    control_names = [
+        name for name ∈ traj_smooth.names
+        if endswith(string(name), string(control_sym))
+    ]
+
+    # Build objective: type-specific infidelity + regularization
+    J = _state_objective(qtraj, traj_smooth, state_sym, Q)
+
+    # Add regularization for control and derivatives
+    J += QuadraticRegularizer(control_names[1], traj_smooth, R_u)
+    J += QuadraticRegularizer(control_names[2], traj_smooth, R_du)
+    J += QuadraticRegularizer(control_names[3], traj_smooth, R_ddu)
+
+    # Add optional Piccolo constraints and objectives
+    J += _apply_piccolo_options(qtraj, piccolo_options, constraints, traj_smooth, state_sym)
+
+    # Initialize dynamics integrators - handle both single integrator and vector of integrators
+    if isnothing(integrator)
+        # Use default BilinearIntegrator for the trajectory type
+        default_int = BilinearIntegrator(qtraj)
+        if default_int isa AbstractVector
+            dynamics_integrators = AbstractIntegrator[default_int...]
+        else
+            dynamics_integrators = AbstractIntegrator[default_int]
+        end
+    elseif integrator isa AbstractIntegrator
+        # Single custom integrator provided
+        dynamics_integrators = AbstractIntegrator[integrator]
+    else
+        # Vector of custom integrators provided
+        dynamics_integrators = AbstractIntegrator[integrator...]
+    end
+
+    # Start with dynamics integrators
+    integrators = copy(dynamics_integrators)
+
+    # Add derivative integrators (always 2 for smooth pulses)
+    push!(integrators, DerivativeIntegrator(control_names[1], control_names[2], traj_smooth))
+
+    push!(integrators, DerivativeIntegrator(control_names[2], control_names[3], traj_smooth))
+
+    if qtraj.system.time_dependent
+        push!(integrators, TimeIntegrator(:t, traj_smooth))
+    end
+
+    prob = DirectTrajOptProblem(
+        traj_smooth,
+        J,
+        integrators;
+        constraints=constraints
+    )
+
+    return QuantumControlProblem(qtraj, prob)
+end
+
+# ============================================================================= #
+# EnsembleTrajectory Constructor
+# ============================================================================= #
+
+@doc raw"""
+    SmoothPulseProblem(qtraj::EnsembleTrajectory; kwargs...)
+
+Construct a `QuantumControlProblem` for smooth pulse optimization over an ensemble of trajectories.
+
+This handles the case where you want to optimize a single pulse that achieves multiple 
+state transfers simultaneously (e.g., |0⟩→|1⟩ and |1⟩→|0⟩ for an X gate via state transfer).
+
+The problem:
+- Creates a combined trajectory with separate state variables for each ensemble member
+- Sums infidelity objectives over all state transfers (weighted by ensemble weights)
+- Creates one dynamics integrator per state variable (all sharing the same controls)
+- Adds derivative integrators for smoothness constraints
+
+# Arguments
+- `qtraj::EnsembleTrajectory`: Ensemble of quantum trajectories with shared system
+
+# Keyword Arguments
+- `integrator::Union{Nothing, AbstractIntegrator, Vector{<:AbstractIntegrator}}=nothing`: Optional custom integrator(s)
+- `du_bound::Float64=Inf`: Bound on first derivative
+- `ddu_bound::Float64=1.0`: Bound on second derivative
+- `Q::Float64=100.0`: Weight on infidelity/objective
+- `R::Float64=1e-2`: Weight on regularization terms (u, u̇, ü)
+- `R_u::Union{Float64, Vector{Float64}}=R`: Weight on control regularization
+- `R_du::Union{Float64, Vector{Float64}}=R`: Weight on first derivative regularization
+- `R_ddu::Union{Float64, Vector{Float64}}=R`: Weight on second derivative regularization
+- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: Additional constraints
+- `piccolo_options::PiccoloOptions=PiccoloOptions()`: Piccolo solver options
+
+# Returns
+- `QuantumControlProblem{EnsembleTrajectory}`: Wrapper containing ensemble trajectory and optimization problem
+
+# Examples
+```julia
+sys = QuantumSystem(H_drift, H_drives, T, drive_bounds)
+
+# Create ensemble for X gate via state transfer
+ψ0 = ComplexF64[1.0, 0.0]
+ψ1 = ComplexF64[0.0, 1.0]
+qtraj1 = KetTrajectory(sys, ψ0, ψ1, N)  # |0⟩ → |1⟩
+qtraj2 = KetTrajectory(sys, ψ1, ψ0, N)  # |1⟩ → |0⟩
+
+ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+qcp = SmoothPulseProblem(ensemble_qtraj; Q=100.0, R=1e-2)
+solve!(qcp; max_iter=100)
+```
+"""
+function SmoothPulseProblem(
+    qtraj::EnsembleTrajectory;
+    integrator::Union{Nothing,AbstractIntegrator,Vector{<:AbstractIntegrator}}=nothing,
+    du_bound::Float64=Inf,
+    ddu_bound::Float64=1.0,
+    Q::Float64=100.0,
+    R::Float64=1e-2,
+    R_u::Union{Float64,Vector{Float64}}=R,
+    R_du::Union{Float64,Vector{Float64}}=R,
+    R_ddu::Union{Float64,Vector{Float64}}=R,
+    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
+    piccolo_options::PiccoloOptions=PiccoloOptions(),
+)
+    if piccolo_options.verbose
+        println("    constructing SmoothPulseProblem for EnsembleTrajectory...")
+        println("\twith $(length(qtraj.trajectories)) trajectories")
+    end
+
+    # Extract info from ensemble trajectory
+    sys = PiccoloQuantumObjects.get_system(qtraj)
+    control_sym = PiccoloQuantumObjects.get_control_name(qtraj)
+    state_names = get_ensemble_state_names(qtraj)
+    weights = get_weights(qtraj)
+    goals = get_goal(qtraj)
+
+    # Build combined trajectory with all ensemble states
+    base_traj, _ = build_ensemble_trajectory_from_trajectories(qtraj.trajectories)
+
+    # Add control derivatives to trajectory
+    du_bounds = fill(du_bound, sys.n_drives)
+    ddu_bounds = fill(ddu_bound, sys.n_drives)
+
+    traj_smooth = add_control_derivatives(
+        base_traj,
+        2;  # Always use 2 derivatives
+        control_name=control_sym,
+        derivative_bounds=(du_bounds, ddu_bounds)
+    )
+
+    # Get control derivative names
+    control_names = [
+        name for name ∈ traj_smooth.names
+        if endswith(string(name), string(control_sym))
+    ]
+
+    # Build objective: weighted sum of infidelities for each state
+    J = _ensemble_state_objective(qtraj, traj_smooth, state_names, weights, goals, Q)
+
+    # Add regularization for control and derivatives
+    J += QuadraticRegularizer(control_names[1], traj_smooth, R_u)
+    J += QuadraticRegularizer(control_names[2], traj_smooth, R_du)
+    J += QuadraticRegularizer(control_names[3], traj_smooth, R_ddu)
+
+    # Apply piccolo options for each state
+    J += _apply_piccolo_options(qtraj, piccolo_options, constraints, traj_smooth, state_names)
+
+    # Build integrators: one dynamics integrator per state
+    if isnothing(integrator)
+        dynamics_integrators = BilinearIntegrator(qtraj, traj_smooth)
+    elseif integrator isa AbstractIntegrator
+        dynamics_integrators = AbstractIntegrator[integrator]
+    else
+        dynamics_integrators = AbstractIntegrator[integrator...]
+    end
+
+    integrators = AbstractIntegrator[dynamics_integrators...]
+
+    # Add derivative integrators (always 2 for smooth pulses)
+    push!(integrators, DerivativeIntegrator(control_names[1], control_names[2], traj_smooth))
+    push!(integrators, DerivativeIntegrator(control_names[2], control_names[3], traj_smooth))
+
+    if sys.time_dependent
+        push!(integrators, TimeIntegrator(:t, traj_smooth))
+    end
+
+    prob = DirectTrajOptProblem(
+        traj_smooth,
+        J,
+        integrators;
+        constraints=constraints
+    )
+
+    return QuantumControlProblem(qtraj, prob)
+end
+
+# ============================================================================= #
+# Type-specific helper functions
+# ============================================================================= #
+
+# ----------------------------------------------------------------------------- #
+# Single trajectory objectives
+# ----------------------------------------------------------------------------- #
+
+# Unitary trajectory: single infidelity objective
+function _state_objective(qtraj::UnitaryTrajectory, traj::NamedTrajectory, state_sym::Symbol, Q::Float64)
+    U_goal = PiccoloQuantumObjects.get_goal(qtraj)
+    return UnitaryInfidelityObjective(U_goal, state_sym, traj; Q=Q)
+end
+
+# Ket trajectory: single infidelity objective
+function _state_objective(qtraj::KetTrajectory, traj::NamedTrajectory, state_sym::Symbol, Q::Float64)
+    return KetInfidelityObjective(state_sym, traj; Q=Q)
+end
+
+# Density trajectory: no fidelity objective yet (TODO)
+function _state_objective(qtraj::DensityTrajectory, traj::NamedTrajectory, state_sym::Symbol, Q::Float64)
+    # TODO: Add fidelity objective when we support general mixed state fidelity
+    return NullObjective(traj)
+end
+
+# Apply Piccolo options with trajectory-type-specific logic
+function _apply_piccolo_options(
+    qtraj::UnitaryTrajectory,
+    piccolo_options::PiccoloOptions,
+    constraints::Vector{<:AbstractConstraint},
+    traj::NamedTrajectory,
+    state_sym::Symbol
+)
+    U_goal = PiccoloQuantumObjects.get_goal(qtraj)
+    return apply_piccolo_options!(
+        piccolo_options, constraints, traj;
+        state_names=state_sym,
+        state_leakage_indices=U_goal isa EmbeddedOperator ?
+                              get_iso_vec_leakage_indices(U_goal) :
+                              nothing
+    )
+end
+
+function _apply_piccolo_options(
+    qtraj::KetTrajectory,
+    piccolo_options::PiccoloOptions,
+    constraints::Vector{<:AbstractConstraint},
+    traj::NamedTrajectory,
+    state_sym::Symbol
+)
+    return apply_piccolo_options!(
+        piccolo_options, constraints, traj;
+        state_names=state_sym
+    )
+end
+
+function _apply_piccolo_options(
+    qtraj::DensityTrajectory,
+    piccolo_options::PiccoloOptions,
+    constraints::Vector{<:AbstractConstraint},
+    traj::NamedTrajectory,
+    state_sym::Symbol
+)
+    return apply_piccolo_options!(
+        piccolo_options, constraints, traj;
+        state_names=state_sym
+    )
+end
+
+# ----------------------------------------------------------------------------- #
+# Ensemble trajectory objectives
+# ----------------------------------------------------------------------------- #
+
+
+"""
+    _ensemble_state_objective(qtraj::EnsembleTrajectory{KetTrajectory}, traj, state_names, weights, goals, Q)
+
+Phase-sensitive ensemble fidelity objective for a goal unitary via state preparations.
+"""
+function _ensemble_state_objective(
+    qtraj::EnsembleTrajectory{KetTrajectory},
+    traj::NamedTrajectory,
+    state_names::Vector{Symbol},
+    weights::Vector{Float64},
+    goals::Vector,
+    Q::Float64
+)
+    # TODO: Is there a use case for weights?
+    return KetsCoherentInfidelityObjective(goals, state_names, traj, Q=Q)
+end
+
+"""
+    _ensemble_state_objective(qtraj::EnsembleTrajectory{UnitaryTrajectory}, traj, state_names, weights, goals, Q)
+
+Create a weighted sum of infidelity objectives for each trajectory in an ensemble.
+"""
+function _ensemble_state_objective(
+    qtraj::EnsembleTrajectory{UnitaryTrajectory},
+    traj::NamedTrajectory,
+    state_names::Vector{Symbol},
+    weights::Vector{Float64},
+    goals::Vector,
+    Q::Float64
+)
+    J = NullObjective(traj)
+    for (name, w, goal) in zip(state_names, weights, goals)
+        J += UnitaryInfidelityObjective(goal, name, traj; Q=w*Q)
+    end
+    return J
+end
+
+function _ensemble_state_objective(
+    qtraj::EnsembleTrajectory{DensityTrajectory},
+    traj::NamedTrajectory,
+    state_names::Vector{Symbol},
+    weights::Vector{Float64},
+    goals::Vector,
+    Q::Float64
+)
+    # TODO: Add fidelity objective when we support general mixed state fidelity
+    return NullObjective(traj)
+end
+
+# ----------------------------------------------------------------------------- #
+# Ensemble piccolo options
+# ----------------------------------------------------------------------------- #
+
+function _apply_piccolo_options(
+    qtraj::EnsembleTrajectory,
+    piccolo_options::PiccoloOptions,
+    constraints::Vector{<:AbstractConstraint},
+    traj::NamedTrajectory,
+    state_names::Vector{Symbol}
+)
+    # Apply piccolo options for all state variables in the ensemble
+    return apply_piccolo_options!(
+        piccolo_options, constraints, traj;
+        state_names=state_names
+    )
+end
+
+# ============================================================================= #
+# Tests
+# ============================================================================= #
+
+@testitem "SmoothPulseProblem with UnitaryTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+    U_goal = GATES[:H]
+    qtraj = UnitaryTrajectory(sys, U_goal, 10)
+    qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+
+    @test qcp isa QuantumControlProblem
+    @test length(qcp.prob.integrators) == 3  # dynamics + du + ddu
+    @test haskey(qcp.prob.trajectory.components, :u)
+    @test haskey(qcp.prob.trajectory.components, :du)
+    @test haskey(qcp.prob.trajectory.components, :ddu)
+
+    # Test accessors
+    @test get_system(qcp) === sys
+    @test get_goal(qcp) === U_goal
+    @test get_trajectory(qcp) === qcp.prob.trajectory
+
+    # Solve and verify
+    solve!(qcp; max_iter=100, print_level=1, verbose=false)
+
+    # Test fidelity after solve
+    traj = get_trajectory(qcp)
+    Ũ⃗_final = traj[end][get_state_name(qtraj)]
+    U_final = iso_vec_to_operator(Ũ⃗_final)
+    fid = unitary_fidelity(U_final, U_goal)
+    @test fid > 0.9
+
+    # Test dynamics constraints are satisfied
+    dynamics_integrator = qcp.prob.integrators[1]
+    δ = zeros(dynamics_integrator.dim)
+    DirectTrajOpt.evaluate!(δ, dynamics_integrator, traj)
+    @test norm(δ, Inf) < 1e-3
+end
+
+@testitem "SmoothPulseProblem with KetTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+    qtraj = KetTrajectory(sys, ψ_init, ψ_goal, 10)
+    qcp = SmoothPulseProblem(qtraj; Q=50.0, R=1e-3)
+
+    @test qcp isa QuantumControlProblem
+    @test length(qcp.prob.integrators) == 3
+    @test haskey(qcp.prob.trajectory.components, :du)
+    @test haskey(qcp.prob.trajectory.components, :ddu)
+
+    # Solve and verify
+    solve!(qcp; max_iter=100, print_level=1, verbose=false)
+
+    # Test fidelity after solve
+    traj = get_trajectory(qcp)
+    ψ̃_final = traj[end][get_state_name(qtraj)]
+    ψ_final = iso_to_ket(ψ̃_final)
+    fid = fidelity(ψ_final, ψ_goal)
+    @test fid > 0.9
+
+    # Test dynamics constraints are satisfied
+    dynamics_integrator = qcp.prob.integrators[1]
+    δ = zeros(dynamics_integrator.dim)
+    DirectTrajOpt.evaluate!(δ, dynamics_integrator, traj)
+    @test norm(δ, Inf) < 1e-3
+end
+
+@testitem "SmoothPulseProblem with DensityTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    sys = OpenQuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+    ρ_init = ComplexF64[1.0 0.0; 0.0 0.0]
+    ρ_goal = ComplexF64[0.0 0.0; 0.0 1.0]
+    qtraj = DensityTrajectory(sys, ρ_init, ρ_goal, 10)
+    qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-3)
+
+    @test qcp isa QuantumControlProblem
+    @test length(qcp.prob.integrators) == 3
+
+    # Solve and verify
+    solve!(qcp; max_iter=100, print_level=1, verbose=false)
+
+    # Test dynamics constraints are satisfied
+    traj = get_trajectory(qcp)
+    dynamics_integrator = qcp.prob.integrators[1]
+    δ = zeros(dynamics_integrator.dim)
+    DirectTrajOpt.evaluate!(δ, dynamics_integrator, traj)
+    @test norm(δ, Inf) < 1e-4
+end
+
+@testitem "SmoothPulseProblem with EnsembleTrajectory (Ket)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    # Use case: Optimize a single pulse that transfers multiple initial states 
+    # to their respective goal states (gate-like behavior for kets)
+    
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+    
+    # Create individual trajectories for different state transfers
+    ψ0 = ComplexF64[1.0, 0.0]
+    ψ1 = ComplexF64[0.0, 1.0]
+    
+    qtraj1 = KetTrajectory(sys, ψ0, ψ1, 10)  # |0⟩ → |1⟩
+    qtraj2 = KetTrajectory(sys, ψ1, ψ0, 10)  # |1⟩ → |0⟩
+    
+    # Create ensemble trajectory
+    ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+    goals = get_goal(ensemble_qtraj)
+    state_names = get_ensemble_state_names(ensemble_qtraj)
+    
+    # Create problem using the new constructor
+    qcp = SmoothPulseProblem(ensemble_qtraj; Q=100.0, R=1e-2)
+
+    @test qcp isa QuantumControlProblem
+    @test qcp.qtraj isa EnsembleTrajectory{KetTrajectory}
+    
+    # Check trajectory components: 2 states + controls + derivatives
+    @test haskey(qcp.prob.trajectory.components, :ψ̃_init_1)
+    @test haskey(qcp.prob.trajectory.components, :ψ̃_init_2)
+    @test haskey(qcp.prob.trajectory.components, :u)
+    @test haskey(qcp.prob.trajectory.components, :du)
+    @test haskey(qcp.prob.trajectory.components, :ddu)
+    
+    # Check integrators: 2 dynamics + 2 derivatives = 4
+    @test length(qcp.prob.integrators) == 4
+
+    # Solve and verify
+    solve!(qcp; max_iter=100, print_level=1, verbose=false)
+
+    # Test fidelity after solve for both states
+    traj = get_trajectory(qcp)
+    for (i, (name, goal)) in enumerate(zip(state_names, goals))
+        ψ̃_final = traj[end][name]
+        ψ_final = iso_to_ket(ψ̃_final)
+        fid = fidelity(ψ_final, goal)
+        @test fid > 0.9
+    end
+
+    # Test dynamics constraints are satisfied for all integrators
+    for integrator in qcp.prob.integrators[1:2]  # First 2 are dynamics
+        δ = zeros(integrator.dim)
+        DirectTrajOpt.evaluate!(δ, integrator, traj)
+        @test norm(δ, Inf) < 1e-4
+    end
+end
+
+@testitem "SmoothPulseProblem with EnsembleTrajectory (Unitary)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    # Use case: Optimize pulses for multiple unitary goals simultaneously
+    # Note: Optimizing for both X and Y gates simultaneously is challenging,
+    # so we use a simpler case: same gate from different initial conditions
+    
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+    
+    # Create trajectories for the same unitary target
+    qtraj1 = UnitaryTrajectory(sys, GATES[:X], 10)
+    qtraj2 = UnitaryTrajectory(sys, GATES[:X], 10)  # Same goal
+    
+    # Create ensemble trajectory
+    ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+    goals = get_goal(ensemble_qtraj)
+    state_names = get_ensemble_state_names(ensemble_qtraj)
+    
+    # Create problem using the new constructor
+    qcp = SmoothPulseProblem(ensemble_qtraj; Q=100.0, R=1e-2)
+
+    @test qcp isa QuantumControlProblem
+    @test qcp.qtraj isa EnsembleTrajectory{UnitaryTrajectory}
+    
+    # Check trajectory components
+    @test haskey(qcp.prob.trajectory.components, :Ũ⃗_init_1)
+    @test haskey(qcp.prob.trajectory.components, :Ũ⃗_init_2)
+    @test haskey(qcp.prob.trajectory.components, :u)
+    @test haskey(qcp.prob.trajectory.components, :du)
+    @test haskey(qcp.prob.trajectory.components, :ddu)
+    
+    # Check integrators: 2 dynamics + 2 derivatives = 4
+    @test length(qcp.prob.integrators) == 4
+
+    # Solve and verify
+    solve!(qcp; max_iter=100, print_level=1, verbose=false)
+
+    # Test fidelity after solve for both unitaries
+    traj = get_trajectory(qcp)
+    for (i, (name, goal)) in enumerate(zip(state_names, goals))
+        Ũ⃗_final = traj[end][name]
+        U_final = iso_vec_to_operator(Ũ⃗_final)
+        fid = unitary_fidelity(U_final, goal)
+        @test fid > 0.9
+    end
+
+    # Test dynamics constraints are satisfied for all integrators
+    for integrator in qcp.prob.integrators[1:2]  # First 2 are dynamics
+        δ = zeros(integrator.dim)
+        DirectTrajOpt.evaluate!(δ, integrator, traj)
+        @test norm(δ, Inf) < 1e-3
+    end
+end
+
+# ============================================================================= #
+# EnsembleTrajectory Tests (multi-state optimization with shared system)
+# ============================================================================= #
+
+@testitem "EnsembleTrajectory multi-state optimization (manual setup)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+
+    # Use case: Optimize a single pulse that transfers multiple initial states 
+    # to their respective goal states (gate-like behavior for kets)
+    
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    
+    # Create individual trajectories for different state transfers
+    ψ0 = ComplexF64[1.0, 0.0]
+    ψ1 = ComplexF64[0.0, 1.0]
+    
+    qtraj1 = KetTrajectory(sys, ψ0, ψ1, 20)  # |0⟩ → |1⟩
+    qtraj2 = KetTrajectory(sys, ψ1, ψ0, 20)  # |1⟩ → |0⟩
+    
+    # Create ensemble trajectory
+    ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+    
+    @test ensemble_qtraj isa EnsembleTrajectory{KetTrajectory}
+    @test get_ensemble_state_names(ensemble_qtraj) == [:ψ̃_init_1, :ψ̃_init_2]
+    
+    # Build combined trajectory with multiple states
+    new_traj, state_names = build_ensemble_trajectory_from_trajectories([qtraj1, qtraj2])
+    
+    @test haskey(new_traj.components, :ψ̃_init_1)
+    @test haskey(new_traj.components, :ψ̃_init_2)
+    @test haskey(new_traj.components, :u)
+    
+    # Verify both initial conditions are set correctly
+    @test new_traj.initial[:ψ̃_init_1] ≈ ket_to_iso(ψ0)
+    @test new_traj.initial[:ψ̃_init_2] ≈ ket_to_iso(ψ1)
+    
+    # Build objective: sum of infidelities for each state
+    goals = get_goal(ensemble_qtraj)
+    J = KetInfidelityObjective(goals[1], state_names[1], new_traj; Q=50.0)
+    J += KetInfidelityObjective(goals[2], state_names[2], new_traj; Q=50.0)
+    
+    # Add regularization
+    J += QuadraticRegularizer(:u, new_traj, 1e-3)
+    
+    # Build integrators: one BilinearIntegrator per state
+    Ĝ = u_ -> sys.G(u_, 0.0)
+    integrators = [
+        BilinearIntegrator(Ĝ, state_names[1], :u, new_traj),
+        BilinearIntegrator(Ĝ, state_names[2], :u, new_traj),
+    ]
+    
+    # Create and solve problem
+    prob = DirectTrajOptProblem(new_traj, J, integrators)
+    qcp = QuantumControlProblem(ensemble_qtraj, prob)
+    
+    solve!(qcp; max_iter=50, verbose=false, print_level=1)
+    
+    # Both state transfers should have reasonable fidelity
+    # (this is essentially implementing an X gate via state transfer)
+end
+
+@testitem "EnsembleTrajectory vs SamplingTrajectory distinction" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using NamedTrajectories
+
+    # This test documents the key difference:
+    # - EnsembleTrajectory: SAME system, DIFFERENT initial/goal states
+    # - SamplingTrajectory: DIFFERENT systems, SAME goal
+    
+    sys = QuantumSystem(GATES[:Z], [GATES[:X]], 1.0, [1.0])
+    
+    # ===== EnsembleTrajectory setup =====
+    # Multiple state transfers on the SAME system
+    ψ0 = ComplexF64[1.0, 0.0]
+    ψ1 = ComplexF64[0.0, 1.0]
+    
+    qtraj_ket1 = KetTrajectory(sys, ψ0, ψ1, 10)
+    qtraj_ket2 = KetTrajectory(sys, ψ1, ψ0, 10)
+    
+    ensemble_qtraj = EnsembleTrajectory([qtraj_ket1, qtraj_ket2])
+    
+    @test get_system(ensemble_qtraj) === sys  # Single system
+    @test length(ensemble_qtraj.trajectories) == 2  # Multiple trajectories
+    @test get_ensemble_state_names(ensemble_qtraj) == [:ψ̃_init_1, :ψ̃_init_2]
+    
+    # ===== SamplingTrajectory setup =====
+    # Same goal, different systems (robust optimization)
+    sys_perturbed = QuantumSystem(1.1 * GATES[:Z], [GATES[:X]], 1.0, [1.0])
+    
+    qtraj_unitary = UnitaryTrajectory(sys, GATES[:X], 10)
+    
+    sampling_qtraj = SamplingTrajectory(qtraj_unitary, [sys, sys_perturbed])
+    
+    @test get_system(sampling_qtraj) === sys  # Nominal system
+    @test length(sampling_qtraj.systems) == 2  # Multiple systems
+    @test get_ensemble_state_names(sampling_qtraj) == [:Ũ⃗_sample_1, :Ũ⃗_sample_2]
+    
+    # Key differences:
+    # 1. EnsembleTrajectory has `trajectories` field (multiple init/goals)
+    # 2. SamplingTrajectory has `systems` field (multiple system params)
+    # 3. State naming: _init_ suffix vs _sample_ suffix
+end
+# ============================================================================= #
+# Time-Dependent System Tests
+# ============================================================================= #
+
+@testitem "SmoothPulseProblem with time-dependent UnitaryTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    # Time-dependent Hamiltonian with oscillating drive
+    ω = 2π * 1.0
+    H(u, t) = GATES[:Z] + u[1] * cos(ω * t) * GATES[:X] + u[2] * sin(ω * t) * GATES[:Y]
+    
+    sys = QuantumSystem(H, 5.0, [1.0, 1.0]; time_dependent=true)
+    
+    U_goal = GATES[:H]
+    qtraj = UnitaryTrajectory(sys, U_goal, 50)
+    
+    # Verify time is in trajectory
+    @test haskey(qtraj.components, :t)
+    @test qtraj.system.time_dependent
+    
+    qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+    
+    @test qcp isa QuantumControlProblem
+    
+    # Should have TimeIntegrator for time-dependent systems
+    has_time_integrator = any(int -> int isa TimeIntegrator, qcp.prob.integrators)
+    @test has_time_integrator
+    
+    # Should have: 1 dynamics + 2 derivative + 1 time = 4 integrators
+    @test length(qcp.prob.integrators) == 4
+    
+    # Solve and verify
+    solve!(qcp; max_iter=50, print_level=5, verbose=false)
+    
+    # Test fidelity after solve
+    traj = get_trajectory(qcp)
+    Ũ⃗_final = traj[end][get_state_name(qtraj)]
+    U_final = iso_vec_to_operator(Ũ⃗_final)
+    fid = unitary_fidelity(U_final, U_goal)
+    @test fid > 0.85
+
+    # Test dynamics constraints are satisfied
+    dynamics_integrator = qcp.prob.integrators[1]
+    δ = zeros(dynamics_integrator.dim)
+    DirectTrajOpt.evaluate!(δ, dynamics_integrator, traj)
+    @test norm(δ, Inf) < 1e-3
+end
+
+@testitem "SmoothPulseProblem with time-dependent KetTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    # Time-dependent Hamiltonian with oscillating drive
+    ω = 2π * 1.0e-1
+    H(u, t) = GATES[:Z] + u[1] * cos(ω * t) * GATES[:X]
+    
+    sys = QuantumSystem(H, 10.0, [1.0]; time_dependent=true)
+    
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+    qtraj = KetTrajectory(sys, ψ_init, ψ_goal, 20)
+    
+    # Verify time is in trajectory
+    @test haskey(qtraj.components, :t)
+    @test qtraj.system.time_dependent
+    
+    qcp = SmoothPulseProblem(qtraj; Q=50.0, R=1e-3)
+    
+    @test qcp isa QuantumControlProblem
+    
+    # Should have TimeIntegrator for time-dependent systems
+    has_time_integrator = any(int -> int isa TimeIntegrator, qcp.prob.integrators)
+    @test has_time_integrator
+    
+    # Should have: 1 dynamics + 2 derivative + 1 time = 4 integrators
+    @test length(qcp.prob.integrators) == 4
+    
+    # Solve and verify
+    solve!(qcp; max_iter=100, print_level=1, verbose=false)
+    
+    # Test fidelity after solve
+    traj = get_trajectory(qcp)
+    ψ̃_final = traj[end][get_state_name(qtraj)]
+    ψ_final = iso_to_ket(ψ̃_final)
+    fid = fidelity(ψ_final, ψ_goal)
+    @test fid > 0.85
+
+    # Test dynamics constraints are satisfied
+    dynamics_integrator = qcp.prob.integrators[1]
+    δ = zeros(dynamics_integrator.dim)
+    DirectTrajOpt.evaluate!(δ, dynamics_integrator, traj)
+    @test norm(δ, Inf) < 1e-3
+end
+
+@testitem "SmoothPulseProblem with time-dependent DensityTrajectory" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    # Time-dependent Hamiltonian with oscillating drive
+    ω = 2π * 5.0
+    H(u, t) = GATES[:Z] + u[1] * cos(ω * t) * GATES[:X] + u[2] * sin(ω * t) * GATES[:Y]
+    
+    sys = OpenQuantumSystem(H, 1.0, [1.0, 1.0]; time_dependent=true)
+    
+    ρ_init = ComplexF64[1.0 0.0; 0.0 0.0]  # |0⟩⟨0|
+    ρ_goal = ComplexF64[0.0 0.0; 0.0 1.0]  # |1⟩⟨1|
+    qtraj = DensityTrajectory(sys, ρ_init, ρ_goal, 20)
+    
+    # Verify time is in trajectory
+    @test haskey(qtraj.components, :t)
+    @test qtraj.system.time_dependent
+    
+    qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-3)
+    
+    @test qcp isa QuantumControlProblem
+    
+    # Should have TimeIntegrator for time-dependent systems
+    has_time_integrator = any(int -> int isa TimeIntegrator, qcp.prob.integrators)
+    @test has_time_integrator
+    
+    # Should have: 1 dynamics + 2 derivative + 1 time = 4 integrators
+    @test length(qcp.prob.integrators) == 4
+    
+    # Solve and verify
+    solve!(qcp; max_iter=50, print_level=1, verbose=false)
+    
+    # Test dynamics constraints are satisfied
+    traj = get_trajectory(qcp)
+    dynamics_integrator = qcp.prob.integrators[1]
+    δ = zeros(dynamics_integrator.dim)
+    DirectTrajOpt.evaluate!(δ, dynamics_integrator, traj)
+    @test norm(δ, Inf) < 1e-4
+end
+
+@testitem "SmoothPulseProblem with time-dependent EnsembleTrajectory (Ket)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    # Time-dependent Hamiltonian with oscillating drive
+    ω = 2π * 5.0
+    H(u, t) = GATES[:Z] + u[1] * cos(ω * t) * GATES[:X] + u[2] * sin(ω * t) * GATES[:Y]
+    
+    sys = QuantumSystem(H, 10.0, [1.0, 1.0]; time_dependent=true)
+    
+    # Create ensemble of ket trajectories for X gate
+    ψ0 = ComplexF64[1.0, 0.0]
+    ψ1 = ComplexF64[0.0, 1.0]
+    
+    qtraj1 = KetTrajectory(sys, ψ0, ψ1, 20)  # |0⟩ → |1⟩
+    qtraj2 = KetTrajectory(sys, ψ1, ψ0, 20)  # |1⟩ → |0⟩
+    
+    # Verify time is in trajectories
+    @test haskey(qtraj1.components, :t)
+    @test haskey(qtraj2.components, :t)
+    
+    ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+    
+    qcp = SmoothPulseProblem(ensemble_qtraj; Q=50.0, R=1e-3)
+    
+    @test qcp isa QuantumControlProblem
+    @test qcp.qtraj isa EnsembleTrajectory{KetTrajectory}
+    
+    # Should have TimeIntegrator for time-dependent systems
+    has_time_integrator = any(int -> int isa TimeIntegrator, qcp.prob.integrators)
+    @test has_time_integrator
+    
+    # Should have: 2 dynamics + 2 derivative + 1 time = 5 integrators
+    @test length(qcp.prob.integrators) == 5
+    
+    # Solve and verify
+    solve!(qcp; max_iter=50, print_level=1, verbose=false)
+    
+    # Test fidelity for both states after solve
+    traj = get_trajectory(qcp)
+    state_names = get_ensemble_state_names(ensemble_qtraj)
+    goals = get_goal(ensemble_qtraj)
+    
+    for (name, goal) in zip(state_names, goals)
+        ψ̃_final = traj[end][name]
+        ψ_final = iso_to_ket(ψ̃_final)
+        fid = fidelity(ψ_final, goal)
+        @test fid > 0.80
+    end
+end
+
+@testitem "SmoothPulseProblem with time-dependent SamplingTrajectory (Unitary)" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+
+    # Time-dependent Hamiltonian with oscillating drive
+    ω = 2π * 5.0
+    H1(u, t) = GATES[:Z] + u[1] * cos(ω * t) * GATES[:X]
+    H2(u, t) = 1.1 * GATES[:Z] + u[1] * cos(ω * t) * GATES[:X]  # Perturbed
+    
+    sys_nominal = QuantumSystem(H1, 1.0, [1.0]; time_dependent=true)
+    sys_perturbed = QuantumSystem(H2, 1.0, [1.0]; time_dependent=true)
+    
+    U_goal = GATES[:X]
+    qtraj = UnitaryTrajectory(sys_nominal, U_goal, 20)
+    
+    # Verify time is in trajectory
+    @test haskey(qtraj.components, :t)
+    
+    qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+    
+    # Create sampling problem
+    sampling_prob = SamplingProblem(qcp, [sys_nominal, sys_perturbed]; Q=100.0)
+    
+    @test sampling_prob isa QuantumControlProblem
+    @test sampling_prob.qtraj isa SamplingTrajectory{UnitaryTrajectory}
+    
+    # Check trajectory has sample states
+    traj = get_trajectory(sampling_prob)
+    @test haskey(traj.components, :Ũ⃗_sample_1)
+    @test haskey(traj.components, :Ũ⃗_sample_2)
+    
+    # Should have TimeIntegrator preserved from base problem
+    has_time_integrator = any(int -> int isa TimeIntegrator, sampling_prob.prob.integrators)
+    @test has_time_integrator
+    
+    # Solve
+    solve!(sampling_prob; max_iter=50, verbose=false, print_level=1)
+    
+    # Test dynamics constraints are satisfied
+    for integrator in sampling_prob.prob.integrators
+        if integrator isa BilinearIntegrator
+            δ = zeros(integrator.dim)
+            DirectTrajOpt.evaluate!(δ, integrator, traj)
+            @test norm(δ, Inf) < 1e-3
+        end
+    end
+end
\ No newline at end of file
diff --git a/src/problem_templates/unitary_free_phase_problem.jl b/src/problem_templates/unitary_free_phase_problem.jl
deleted file mode 100644
index 271fd144..00000000
--- a/src/problem_templates/unitary_free_phase_problem.jl
+++ /dev/null
@@ -1,162 +0,0 @@
-export UnitaryFreePhaseProblem
-
-
-"""
-    UnitaryFreePhaseProblem(system::AbstractQuantumSystem, goal::Function, T, Δt; kwargs...)
-
-
-Construct a `DirectTrajOptProblem` for a free-time unitary gate problem with smooth control pulses enforced by constraining the second derivative of the pulse trajectory. The problem follows the same structure as `UnitarySmoothPulseProblem`, but allows for free global phases on the goal unitary, via cosines and sines parameterizing phase variables.
-    
-The `goal` function should accept a vector of global phases `[cos(θ); sin(θ)]` and return an `AbstractPiccoloOperator`.
-"""
-function UnitaryFreePhaseProblem(
-    system::AbstractQuantumSystem,
-    goal::Function,
-    T::Int,
-    Δt::Union{Float64, AbstractVector{Float64}};
-    unitary_integrator=UnitaryIntegrator,
-    state_name::Symbol = :Ũ⃗,
-    control_name::Symbol = :u,
-    timestep_name::Symbol = :Δt,
-    phase_name::Symbol = :θ,
-    init_phases::Union{AbstractVector{Float64}, Nothing}=nothing,
-    init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
-    u_guess::Union{Matrix{Float64}, Nothing}=nothing,
-    du_bound::Float64=Inf,
-    du_bounds::Vector{Float64}=fill(du_bound, system.n_drives),
-    ddu_bound::Float64=1.0,
-    ddu_bounds::Vector{Float64}=fill(ddu_bound, system.n_drives),
-    Δt_min::Float64=0.5 * minimum(Δt),
-    Δt_max::Float64=2.0 * maximum(Δt),
-    Q::Float64=100.0,
-    R=1e-2,
-    R_u::Union{Float64, Vector{Float64}}=R,
-    R_du::Union{Float64, Vector{Float64}}=R,
-    R_ddu::Union{Float64, Vector{Float64}}=R,
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-)
-    if piccolo_options.verbose
-        println("    constructing UnitaryFreePhaseProblem...")
-        println("\tusing integrator: $(typeof(unitary_integrator))")
-        println("\tinitial free phases: $(phase_name) = $(init_phases)")
-    end
-
-    # Construct global phases
-    x = Symbol("cos$(phase_name)")
-    y = Symbol("sin$(phase_name)")
-    phase_names = [x, y]
-    if piccolo_options.verbose
-        println("\tusing global names: ", phase_names)
-    end
-
-    # Trajectory
-    if !isnothing(init_trajectory)
-        trig_phases = [init_trajectory.global_data[x]; init_trajectory.global_data[y]]
-        @assert goal(trig_phases) isa AbstractPiccoloOperator "expected goal([cos(θ); sin(θ)])"
-        eval_goal = goal(trig_phases)
-        traj = init_trajectory
-    else
-        phase_data = (; x => cos.(init_phases), y => sin.(init_phases))
-        trig_phases = [phase_data[x]; phase_data[y]]
-        @assert goal(trig_phases) isa AbstractPiccoloOperator "expected goal([cos(θ); sin(θ)])"
-        eval_goal = goal(trig_phases)
-        traj = initialize_trajectory(
-            eval_goal,
-            T,
-            Δt,
-            system.n_drives,
-            (system.drive_bounds, du_bounds, ddu_bounds);
-            state_name=state_name,
-            control_name=control_name,
-            timestep_name=timestep_name,
-            Δt_bounds=(Δt_min, Δt_max),
-            zero_initial_and_final_derivative=piccolo_options.zero_initial_and_final_derivative,
-            geodesic=piccolo_options.geodesic,
-            bound_state=piccolo_options.bound_state,
-            u_guess=u_guess,
-            system=system,
-            rollout_integrator=piccolo_options.rollout_integrator,
-            verbose=piccolo_options.verbose,
-            global_component_data=phase_data,
-        )
-    end
-
-    # Objective
-    J = UnitaryFreePhaseInfidelityObjective(goal,  state_name,  phase_names, traj; Q=Q)
-
-    control_names = [
-        name for name ∈ traj.names
-            if endswith(string(name), string(control_name))
-    ]
-
-    J += QuadraticRegularizer(control_names[1], traj, R_u)
-    J += QuadraticRegularizer(control_names[2], traj, R_du)
-    J += QuadraticRegularizer(control_names[3], traj, R_ddu)
-
-    # Optional Piccolo constraints and objectives
-    J += apply_piccolo_options!(
-        piccolo_options, constraints, traj;
-        state_names=state_name,
-        state_leakage_indices=eval_goal isa EmbeddedOperator ? 
-            get_iso_vec_leakage_indices(eval_goal) :
-            nothing
-    )
-    
-    # Phase constraint
-    function phase_norm(z)
-        x, y = z[1:length(z) ÷ 2], z[length(z) ÷ 2 + 1:end]
-        return x .^ 2 + y .^2 .- 1
-    end
-    push!(constraints, NonlinearGlobalConstraint(phase_norm, phase_names, traj))
-
-    integrators = [
-        unitary_integrator(system, traj, state_name, control_name),
-        DerivativeIntegrator(traj, control_name, control_names[2]),
-        DerivativeIntegrator(traj, control_names[2], control_names[3]),
-    ]
-
-    return DirectTrajOptProblem(
-        traj,
-        J,
-        integrators;
-        constraints=constraints
-    )
-end
-
-# *************************************************************************** #
-
-@testitem "UnitaryFreePhaseProblem: basic construction" begin
-    using PiccoloQuantumObjects
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0)]
-
-    sys = QuantumSystem(0.3 * PAULIS.X, [PAULIS.Y], T_max, u_bounds)
-    U_goal = GATES.Z
-    T = 51
-    Δt = 0.2
-
-    function virtual_z(z::AbstractVector{<:Real})        
-        x, y = z[1:length(z)÷2], z[1+length(z)÷2:end]
-        # U_goal ≈ R * U
-        R = reduce(kron, [xᵢ * PAULIS.I + im * yᵢ * PAULIS.Z for (xᵢ, yᵢ) in zip(x, y)])
-        return R'U_goal
-    end
-
-    initial_phases = [pi/3]
-
-    prob = UnitaryFreePhaseProblem(
-        sys, virtual_z, T, Δt, 
-        init_phases=initial_phases,
-        piccolo_options=PiccoloOptions(verbose=false),
-        phase_name=:ϕ,
-    )
-
-    @test prob isa DirectTrajOptProblem
-    @test length(prob.trajectory.global_data) == 2length(initial_phases)
-    @test prob.trajectory.global_names == (:cosϕ, :sinϕ)
-
-    before = copy(prob.trajectory.global_data)
-    solve!(prob, max_iter=10, verbose=false, print_level=1)
-    @test prob.trajectory.global_data != before
-end
diff --git a/src/problem_templates/unitary_minimum_time_problem.jl b/src/problem_templates/unitary_minimum_time_problem.jl
deleted file mode 100644
index d77c1865..00000000
--- a/src/problem_templates/unitary_minimum_time_problem.jl
+++ /dev/null
@@ -1,196 +0,0 @@
-export UnitaryMinimumTimeProblem
-
-
-@doc raw"""
-    UnitaryMinimumTimeProblem(
-        goal::AbstractPiccoloOperator,
-        trajectory::NamedTrajectory,
-        objective::Objective,
-        dynamics::TrajectoryDynamics,
-        constraints::AbstractVector{<:AbstractConstraint};
-        kwargs...
-    )
-
-    UnitaryMinimumTimeProblem(
-        goal::AbstractPiccoloOperator,
-        prob::DirectTrajOptProblem;
-        kwargs...
-    )
-
-Create a minimum-time problem for unitary control.
-
-```math
-\begin{aligned}
-\underset{\vec{\tilde{U}}, u, \dot{u}, \ddot{u}, \Delta t}{\text{minimize}} & \quad
-J(\vec{\tilde{U}}, u, \dot{u}, \ddot{u}) + D \sum_t \Delta t_t \\
-\text{ subject to } & \quad \vb{P}^{(n)}\qty(\vec{\tilde{U}}_{t+1}, \vec{\tilde{U}}_t, u_t, \Delta t_t) = 0 \\
-& c(\vec{\tilde{U}}, u, \dot{u}, \ddot{u}) = 0 \\
-& \quad \Delta t_{\text{min}} \leq \Delta t_t \leq \Delta t_{\text{max}} \\
-\end{aligned}
-```
-
-# Keyword Arguments
-- `piccolo_options::PiccoloOptions=PiccoloOptions()`: The Piccolo options.
-- `unitary_name::Symbol=:Ũ⃗`: The name of the unitary for the goal.
-- `final_fidelity::Float64=1.0`: The final fidelity constraint.
-- `D::Float64=1.0`: The scaling factor for the minimum-time objective.
-"""
-function UnitaryMinimumTimeProblem end
-
-function UnitaryMinimumTimeProblem(
-    trajectory::NamedTrajectory,
-    goal::AbstractPiccoloOperator,
-    objective::Objective,
-    dynamics::TrajectoryDynamics,
-    constraints::AbstractVector{<:AbstractConstraint};
-    unitary_name::Symbol=:Ũ⃗,
-    final_fidelity::Float64=1.0,
-    D::Float64=100.0,
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-)
-    if piccolo_options.verbose
-        println("    constructing UnitaryMinimumTimeProblem...")
-        println("\tfinal fidelity: $(final_fidelity)")
-    end
-
-    objective += MinimumTimeObjective(trajectory, D=D)
-
-    fidelity_constraint = FinalUnitaryFidelityConstraint(
-        goal, unitary_name, final_fidelity, trajectory
-    )
-
-    constraints = push!(constraints, fidelity_constraint)
-
-    return DirectTrajOptProblem(
-        trajectory,
-        objective,
-        dynamics,
-        constraints
-    )
-end
-
-# TODO: how to handle apply_piccolo_options?
-# TODO: goal from trajectory?
-
-function UnitaryMinimumTimeProblem(
-    prob::DirectTrajOptProblem,
-    goal::AbstractPiccoloOperator;
-    objective::Objective=prob.objective,
-    constraints::AbstractVector{<:AbstractConstraint}=deepcopy(prob.constraints),
-    kwargs...
-)
-    return UnitaryMinimumTimeProblem(
-        deepcopy(prob.trajectory),
-        goal,
-        objective,
-        prob.dynamics,
-        constraints;
-        kwargs...
-    )
-end
-
-# --------------------------------------------------------------------------- #
-# Free phases
-# --------------------------------------------------------------------------- #
-
-function UnitaryMinimumTimeProblem(
-    trajectory::NamedTrajectory,
-    goal::Function,
-    objective::Objective,
-    dynamics::TrajectoryDynamics,
-    constraints::AbstractVector{<:AbstractConstraint};
-    unitary_name::Symbol=:Ũ⃗,
-    phase_name::Symbol=:θ,
-    final_fidelity::Float64=1.0,
-    D::Float64=100.0,
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-)
-    # Collect phase names
-    phase_names = [n for n in trajectory.global_names if endswith(string(n), string(phase_name))]
-
-    if piccolo_options.verbose
-        println("    constructing UnitaryMinimumTimeProblem...")
-        println("\tfinal fidelity: $(final_fidelity)")
-        println("\tphase names: $(phase_names)")
-    end
-
-    objective += MinimumTimeObjective(trajectory, D=D)
-    # timesteps_all_equal=piccolo_options.timesteps_all_equal
-
-    fidelity_constraint = FinalUnitaryFidelityConstraint(
-        goal, unitary_name, phase_names, final_fidelity, trajectory
-    )
-
-    constraints = push!(constraints, fidelity_constraint)
-
-    return DirectTrajOptProblem(
-        trajectory,
-        objective,
-        dynamics,
-        constraints
-    )
-end
-
-function UnitaryMinimumTimeProblem(
-    prob::DirectTrajOptProblem,
-    goal::Function;
-    objective::Objective=prob.objective,
-    constraints::AbstractVector{<:AbstractConstraint}=deepcopy(prob.constraints),
-    kwargs...
-)
-    return UnitaryMinimumTimeProblem(
-        deepcopy(prob.trajectory),
-        goal,
-        objective,
-        deepcopy(prob.dynamics),
-        constraints;
-        kwargs...
-    )
-end
-
-# *************************************************************************** #
-
-@testitem "Minimum time Hadamard gate" begin
-    using NamedTrajectories
-    using PiccoloQuantumObjects 
-
-    H_drift = 0.1PAULIS[:Z]
-    H_drives = [PAULIS[:X], PAULIS[:Y]]
-    U_goal = GATES[:H]
-    T = 51
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    Δt = 0.2
-
-    sys = QuantumSystem(H_drift, H_drives, T_max, u_bounds)
-
-    prob = UnitarySmoothPulseProblem(
-        sys, U_goal, T, Δt, Δt_min=Δt * 0.01,
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-
-    before = unitary_rollout_fidelity(prob.trajectory, sys, drive_name=:u)
-    solve!(prob; max_iter=150, verbose=false, print_level=1)
-    after = unitary_rollout_fidelity(prob.trajectory, sys, drive_name=:u)
-    @test after > before
-
-    # soft fidelity constraint
-    min_prob = UnitaryMinimumTimeProblem(
-        prob, U_goal,
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    solve!(min_prob; max_iter=150, verbose=false, print_level=1)
-
-    # test fidelity has stayed above the constraint
-    constraint_tol = 0.95
-    final_fidelity = minimum([0.99, after])
-    @test unitary_rollout_fidelity(min_prob.trajectory, sys, drive_name=:u) ≥ constraint_tol * final_fidelity
-    duration_after = sum(get_timesteps(min_prob.trajectory))
-    duration_before = sum(get_timesteps(prob.trajectory))
-    @test duration_after <= duration_before
-end
-
-@testitem "Test relaxed final_fidelity constraint" begin
-    final_fidelity = 0.95
-    @test_broken false
-end
diff --git a/src/problem_templates/unitary_sampling_problem.jl b/src/problem_templates/unitary_sampling_problem.jl
deleted file mode 100644
index 1578a65e..00000000
--- a/src/problem_templates/unitary_sampling_problem.jl
+++ /dev/null
@@ -1,200 +0,0 @@
-export UnitarySamplingProblem
-
-
-@doc raw"""
-    UnitarySamplingProblem(systemns, operator, T, Δt; kwargs...)
-
-A `UnitarySamplingProblem` is a quantum control problem where the goal is to find a control
-pulse that generates a target unitary operator for a set of quantum systems.
-The controls are shared among all systems, and the optimization seeks to find the control
-pulse that achieves the operator for each system. The idea is to enforce a
-robust solution by including multiple systems reflecting the problem uncertainty.
-
-# Arguments
-- `systems::AbstractVector{<:AbstractQuantumSystem}`: A vector of quantum systems.
-- `operators::AbstractVector{<:AbstractPiccoloOperator}`: A vector of target operators.
-- `T::Int`: The number of time steps.
-- `Δt::Union{Float64, Vector{Float64}}`: The time step value or vector of time steps.
-
-# Keyword Arguments
-- `system_labels::Vector{String} = string.(1:length(systems))`: The labels for each system.
-- `system_weights::Vector{Float64} = fill(1.0, length(systems))`: The weights for each system.
-- `init_trajectory::Union{NamedTrajectory, Nothing} = nothing`: The initial trajectory.
-- `state_name::Symbol = :Ũ⃗`: The name of the state variable.
-- `control_name::Symbol = :u`: The name of the control variable.
-- `timestep_name::Symbol = :Δt`: The name of the timestep variable.
-- `constraints::Vector{<:AbstractConstraint} = AbstractConstraint[]`: The constraints.
-- `u_bound::Float64 = 1.0`: The bound for the control amplitudes.
-- `u_bounds = fill(u_bound, length(systems[1].G_drives))`: The bounds for the control amplitudes.
-- `u_guess::Union{Matrix{Float64}, Nothing} = nothing`: The initial guess for the control amplitudes.
-- `du_bound::Float64 = Inf`: The bound for the control first derivatives.
-- `du_bounds = fill(du_bound, length(systems[1].G_drives))`: The bounds for the control first derivatives.
-- `ddu_bound::Float64 = 1.0`: The bound for the control second derivatives.
-- `ddu_bounds = fill(ddu_bound, length(systems[1].G_drives))`: The bounds for the control second derivatives.
-- `Δt_min::Float64 = 0.5 * Δt`: The minimum time step size.
-- `Δt_max::Float64 = 1.5 * Δt`: The maximum time step size.
-- `Q::Float64 = 100.0`: The fidelity weight.
-- `R::Float64 = 1e-2`: The regularization weight.
-- `R_u::Union{Float64, Vector{Float64}} = R`: The regularization weight for the control amplitudes.
-- `R_du::Union{Float64, Vector{Float64}} = R`: The regularization weight for the control first derivatives.
-- `R_ddu::Union{Float64, Vector{Float64}} = R`: The regularization weight for the control second derivatives.
-- `piccolo_options::PiccoloOptions = PiccoloOptions()`: The Piccolo options.
-
-"""
-function UnitarySamplingProblem(
-    systems::AbstractVector{<:AbstractQuantumSystem},
-    operators::AbstractVector{<:AbstractPiccoloOperator},
-    T::Int,
-    Δt::Union{Float64,Vector{Float64}};
-    unitary_integrator=UnitaryIntegrator,
-    system_weights=fill(1.0, length(systems)),
-    init_trajectory::Union{NamedTrajectory,Nothing}=nothing,
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-    state_name::Symbol=:Ũ⃗,
-    control_name::Symbol=:u,
-    timestep_name::Symbol=:Δt,
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
-    u_guess::Union{Matrix{Float64}, Nothing}=nothing,
-    du_bound::Float64=Inf,
-    du_bounds=fill(du_bound, systems[1].n_drives),
-    ddu_bound::Float64=1.0,
-    ddu_bounds=fill(ddu_bound, systems[1].n_drives),
-    Δt_min::Float64=0.5 * minimum(Δt),
-    Δt_max::Float64=2.0 * maximum(Δt),
-    Q::Float64=100.0,
-    R=1e-2,
-    R_u::Union{Float64, Vector{Float64}}=R,
-    R_du::Union{Float64, Vector{Float64}}=R,
-    R_ddu::Union{Float64, Vector{Float64}}=R,
-    kwargs...
-)
-    @assert length(systems) == length(operators)
-
-    if piccolo_options.verbose
-        println("    constructing UnitarySamplingProblem...")
-        println("\tusing integrator: $(typeof(unitary_integrator))")
-        println("\tusing $(length(systems)) systems")
-    end
-
-    # Trajectory
-    state_names = [
-        Symbol(string(state_name, "_system_", label)) for label ∈ 1:length(systems)
-    ]
-
-    if !isnothing(init_trajectory)
-        traj = init_trajectory
-    else
-        trajs = map(zip(systems, operators, state_names)) do (sys, op, st)
-            initialize_trajectory(
-                op,
-                T,
-                Δt,
-                sys.n_drives,
-                (sys.drive_bounds, du_bounds, ddu_bounds);
-                state_name=st,
-                control_name=control_name,
-                timestep_name=timestep_name,
-                Δt_bounds=(Δt_min, Δt_max),
-                geodesic=piccolo_options.geodesic,
-                bound_state=piccolo_options.bound_state,
-                u_guess=u_guess,
-                system=sys,
-                rollout_integrator=piccolo_options.rollout_integrator,
-                verbose=false # loop
-            )
-        end
-
-        traj = merge(
-            trajs, merge_names=(u=1, du=1, ddu=1, Δt=1), timestep=timestep_name
-        )
-    end
-
-    control_names = [
-        name for name ∈ traj.names
-        if endswith(string(name), string(control_name))
-    ]
-
-    # Objective
-    J = QuadraticRegularizer(control_name, traj, R_u)
-    J += QuadraticRegularizer(control_names[2], traj, R_du)
-    J += QuadraticRegularizer(control_names[3], traj, R_ddu)
-
-    for (weight, op, name) in zip(system_weights, operators, state_names)
-        J += UnitaryInfidelityObjective(op, name, traj; Q=weight * Q)
-    end
-
-    # Optional Piccolo constraints and objectives
-    J += apply_piccolo_options!(
-        piccolo_options, constraints, traj;
-        state_names=state_names,
-        state_leakage_indices=all(op -> op isa EmbeddedOperator, operators) ?       
-            get_iso_vec_leakage_indices.(operators) : 
-            nothing
-    )
-
-    # Integrators
-    integrators = [
-        [unitary_integrator(sys, traj, n, control_name) for (sys, n) in zip(systems, state_names)]...,
-        DerivativeIntegrator(traj, control_name, control_names[2]),
-        DerivativeIntegrator(traj, control_names[2], control_names[3]),
-    ]
-
-    return DirectTrajOptProblem(
-        traj,
-        J,
-        integrators;
-        constraints=constraints,
-    )
-end
-
-function UnitarySamplingProblem(
-    systems::AbstractVector{<:AbstractQuantumSystem},
-    operator::AbstractPiccoloOperator,
-    T::Int,
-    Δt::Union{Float64,Vector{Float64}};
-    kwargs...
-)
-    # Broadcast the operator to all systems
-    return UnitarySamplingProblem(
-        systems,
-        fill(operator, length(systems)),
-        T,
-        Δt;
-        kwargs...
-    )
-end
-
-# *************************************************************************** #
-
-@testitem "Sample robustness test" begin
-    using PiccoloQuantumObjects
-
-    T = 50
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    Δt = 0.2
-    timesteps = fill(Δt, T)
-    operator = GATES[:H]
-    systems(ζ) = QuantumSystem(ζ * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-
-    samples = [0.0, 0.1]
-    prob = UnitarySamplingProblem(
-        [systems(x) for x in samples], operator, T, Δt,
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    solve!(prob, max_iter=100, print_level=1, verbose=false)
-    
-    base_prob = UnitarySmoothPulseProblem(
-        systems(samples[1]), operator, T, Δt,
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    solve!(base_prob, max_iter=100, verbose=false, print_level=1)
-    
-    fid = []
-    base_fid = []
-    for x in range(samples[1], samples[1], length=5)
-        push!(fid, unitary_rollout_fidelity(prob.trajectory, systems(0.1), unitary_name=:Ũ⃗_system_1, drive_name=:u))
-        push!(base_fid, unitary_rollout_fidelity(base_prob.trajectory, systems(0.1), drive_name=:u))
-    end
-    @test sum(fid) > sum(base_fid)
-end
diff --git a/src/problem_templates/unitary_smooth_pulse_problem.jl b/src/problem_templates/unitary_smooth_pulse_problem.jl
deleted file mode 100644
index 4536c24e..00000000
--- a/src/problem_templates/unitary_smooth_pulse_problem.jl
+++ /dev/null
@@ -1,257 +0,0 @@
-export UnitarySmoothPulseProblem
-
-
-@doc raw"""
-    UnitarySmoothPulseProblem(system::AbstractQuantumSystem, operator::AbstractPiccoloOperator, T::Int, Δt::Float64; kwargs...)
-    UnitarySmoothPulseProblem(H_drift, H_drives, operator, T, Δt; kwargs...)
-
-Construct a `DirectTrajOptProblem` for a free-time unitary gate problem with smooth control pulses enforced by constraining the second derivative of the pulse trajectory, i.e.,
-
-```math
-\begin{aligned}
-\underset{\vec{\tilde{U}}, u, \dot{u}, \ddot{u}, \Delta t}{\text{minimize}} & \quad
-Q \cdot \ell\qty(\vec{\tilde{U}}_T, \vec{\tilde{U}}_{\text{goal}}) + \frac{1}{2} \sum_t \qty(R_u u_t^2 + R_{\dot{u}} \dot{u}_t^2 + R_{\ddot{u}} \ddot{u}_t^2) \\
-\text{ subject to } & \quad \vb{P}^{(n)}\qty(\vec{\tilde{U}}_{t+1}, \vec{\tilde{U}}_t, u_t, \Delta t_t) = 0 \\
-& \quad u_{t+1} - u_t - \dot{u}_t \Delta t_t = 0 \\
-& \quad \dot{u}_{t+1} - \dot{u}_t - \ddot{u}_t \Delta t_t = 0 \\
-& \quad |u_t| \leq u_{\text{bound}} \\
-& \quad |\ddot{u}_t| \leq \ddot{u}_{\text{bound}} \\
-& \quad \Delta t_{\text{min}} \leq \Delta t_t \leq \Delta t_{\text{max}} \\
-\end{aligned}
-```
-
-where, for $U \in SU(N)$,
-
-```math
-\ell\qty(\vec{\tilde{U}}_T, \vec{\tilde{U}}_{\text{goal}}) =
-\abs{1 - \frac{1}{N} \abs{ \tr \qty(U_{\text{goal}}, U_T)} }
-```
-
-is the *infidelity* objective function, $Q$ is a weight, $R_u$, $R_{\dot{u}}$, and $R_{\ddot{u}}$ are weights on the regularization terms, and $\vb{P}^{(n)}$ is the $n$th-order Pade integrator.
-
-# Arguments
-
-- `system::AbstractQuantumSystem`: the system to be controlled
-or
-- `H_drift::AbstractMatrix{<:Number}`: the drift hamiltonian
-- `H_drives::Vector{<:AbstractMatrix{<:Number}}`: the control hamiltonians
-with
-- `goal::AbstractPiccoloOperator`: the target unitary, either in the form of an `EmbeddedOperator` or a `Matrix{ComplexF64}
-- `T::Int`: the number of timesteps
-- `Δt::Float64`: the (initial) time step size
-
-# Keyword Arguments
-- `piccolo_options::PiccoloOptions=PiccoloOptions()`: the options for the Piccolo solver
-- `state_name::Symbol = :Ũ⃗`: the name of the state
-- `control_name::Symbol = :u`: the name of the control
-- `timestep_name::Symbol = :Δt`: the name of the timestep
-- `init_trajectory::Union{NamedTrajectory, Nothing}=nothing`: an initial trajectory to use
-- `u_guess::Union{Matrix{Float64}, Nothing}=nothing`: an initial guess for the control pulses
-- `u_bound::Float64=1.0`: the bound on the control pulse
-- `u_bounds=fill(u_bound, length(system.G_drives))`: the bounds on the control pulses, one for each drive
-- `du_bound::Float64=Inf`: the bound on the control pulse derivative
-- `du_bounds=fill(du_bound, length(system.G_drives))`: the bounds on the control pulse derivatives, one for each drive
-- `ddu_bound::Float64=1.0`: the bound on the control pulse second derivative
-- `ddu_bounds=fill(ddu_bound, length(system.G_drives))`: the bounds on the control pulse second derivatives, one for each drive
-- `Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * mean(Δt)`: the minimum time step size
-- `Δt_max::Float64=Δt isa Float64 ? 1.5 * Δt : 1.5 * mean(Δt)`: the maximum time step size
-- `Q::Float64=100.0`: the weight on the infidelity objective
-- `R=1e-2`: the weight on the regularization terms
-- `R_u::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulses
-- `R_du::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulse derivatives
-- `R_ddu::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulse second derivatives
-- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: the constraints to enforce
-
-"""
-function UnitarySmoothPulseProblem end
-
-function UnitarySmoothPulseProblem(
-    system::AbstractQuantumSystem,
-    goal::AbstractPiccoloOperator,
-    N::Int,
-    Δt::Union{Float64, <:AbstractVector{Float64}};
-    unitary_integrator=UnitaryIntegrator,
-    state_name::Symbol = :Ũ⃗,
-    control_name::Symbol = :u,
-    timestep_name::Symbol = :Δt,
-    init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
-    u_guess::Union{Matrix{Float64}, Nothing}=nothing,
-    du_bound::Float64=Inf,
-    du_bounds=fill(du_bound, system.n_drives),
-    ddu_bound::Float64=1.0,
-    ddu_bounds=fill(ddu_bound, system.n_drives),
-    Δt_min::Float64=0.5 * minimum(Δt),
-    Δt_max::Float64=2.0 * maximum(Δt),
-    Q::Float64=100.0,
-    R=1e-2,
-    R_u::Union{Float64, Vector{Float64}}=R,
-    R_du::Union{Float64, Vector{Float64}}=R,
-    R_ddu::Union{Float64, Vector{Float64}}=R,
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-)
-    if piccolo_options.verbose
-        println("    constructing UnitarySmoothPulseProblem...")
-        println("\tusing integrator: $(typeof(unitary_integrator))")
-    end
-
-    # Trajectory
-    if !isnothing(init_trajectory)
-        traj = init_trajectory
-    else
-        u_bounds = (
-            [system.drive_bounds[j][1] for j in 1:length(system.drive_bounds)], 
-            [system.drive_bounds[j][2] for j in 1:length(system.drive_bounds)]
-        )
-
-        traj = initialize_trajectory(
-            goal,
-            N,
-            Δt,
-            system.n_drives,
-            (u_bounds, du_bounds, ddu_bounds);
-            state_name=state_name,
-            control_name=control_name,
-            timestep_name=timestep_name,
-            Δt_bounds=(Δt_min, Δt_max),
-            zero_initial_and_final_derivative=piccolo_options.zero_initial_and_final_derivative,
-            geodesic=piccolo_options.geodesic,
-            bound_state=piccolo_options.bound_state,
-            u_guess=u_guess,
-            system=system,
-            rollout_integrator=piccolo_options.rollout_integrator,
-            verbose=piccolo_options.verbose
-        )
-    end
-
-    # Objective
-    J = UnitaryInfidelityObjective(goal, state_name, traj; Q=Q)
-
-    control_names = [
-        name for name ∈ traj.names
-            if endswith(string(name), string(control_name))
-    ]
-
-    J += QuadraticRegularizer(control_names[1], traj, R_u)
-    J += QuadraticRegularizer(control_names[2], traj, R_du)
-    J += QuadraticRegularizer(control_names[3], traj, R_ddu)
-
-    # Optional Piccolo constraints and objectives
-    J += apply_piccolo_options!(
-        piccolo_options, constraints, traj; 
-        state_names=state_name,
-        state_leakage_indices=goal isa EmbeddedOperator ? 
-            get_iso_vec_leakage_indices(goal) : 
-            nothing
-    )
-
-    integrators = [
-        unitary_integrator(system, traj, state_name, control_name),
-        DerivativeIntegrator(traj, control_name, control_names[2]),
-        DerivativeIntegrator(traj, control_names[2], control_names[3]),
-    ]
-
-    return DirectTrajOptProblem(
-        traj,
-        J,
-        integrators;
-        constraints=constraints
-    )
-end
-
-function UnitarySmoothPulseProblem(
-    H_drift::AbstractMatrix{<:Number},
-    H_drives::Vector{<:AbstractMatrix{<:Number}},
-    u_bounds::PiccoloQuantumObjects.QuantumSystems.DriveBounds,
-    args...;
-    kwargs...
-)
-    system = QuantumSystem(H_drift, H_drives, T_max, u_bounds)
-    return UnitarySmoothPulseProblem(system, args...; kwargs...)
-end
-
-# *************************************************************************** #
-
-@testitem "Hadamard gate improvement" begin
-    using PiccoloQuantumObjects
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-    U_goal = GATES[:H]
-    T = 51
-    Δt = 0.2
-
-    prob = UnitarySmoothPulseProblem(
-        sys, U_goal, T, Δt;
-        du_bound=1.0,
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-
-    initial = unitary_rollout_fidelity(prob.trajectory, sys, drive_name=:u)
-    solve!(prob, max_iter=100, verbose=false, print_level=1)
-    @test unitary_rollout_fidelity(prob.trajectory, sys, drive_name=:u) > initial
-end
-
-@testitem "Bound states and control norm constraint" begin
-    using PiccoloQuantumObjects
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
-    U_goal = GATES[:H]
-    T = 51
-    Δt = 0.2
-
-    prob = UnitarySmoothPulseProblem(
-        sys, U_goal, T, Δt,
-        piccolo_options=PiccoloOptions(
-            verbose=false,
-            bound_state=true,
-            complex_control_norm_constraint_name=:u
-        )
-    )
-
-    initial = copy(prob.trajectory.data)
-    solve!(prob, max_iter=5, verbose=false, print_level=1)
-    @test prob.trajectory.data != initial
-end
-
-@testitem "EmbeddedOperator tests" begin
-    using PiccoloQuantumObjects 
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-
-    a = annihilate(3)
-    sys = QuantumSystem([(a + a')/2, (a - a')/(2im)], T_max, u_bounds)
-    U_goal = EmbeddedOperator(GATES[:H], sys)
-    T = 51
-    Δt = 0.2
-
-    @testset "EmbeddedOperator: solve gate" begin
-        prob = UnitarySmoothPulseProblem(
-            sys, U_goal, T, Δt,
-            piccolo_options=PiccoloOptions(verbose=false)
-        )
-
-        initial = copy(prob.trajectory.data)
-        solve!(prob, max_iter=5, verbose=false, print_level=1)
-        @test prob.trajectory.data != initial
-    end
-
-    @testset "EmbeddedOperator: leakage constraint" begin
-        prob = UnitarySmoothPulseProblem(
-            sys, U_goal, T, Δt;
-            du_bound=1.0,
-            piccolo_options=PiccoloOptions(
-                leakage_constraint=true, 
-                leakage_constraint_value=5e-2, 
-                leakage_cost=1e-1,
-                verbose=false
-            )
-        )
-        initial = copy(prob.trajectory.data)
-        solve!(prob, max_iter=5, verbose=false, print_level=1)
-        @test prob.trajectory.data != initial
-    end
-end
-
-# TODO: Test changing names of control, state, and timestep
diff --git a/src/problem_templates/unitary_variational_problem.jl b/src/problem_templates/unitary_variational_problem.jl
deleted file mode 100644
index 0253188d..00000000
--- a/src/problem_templates/unitary_variational_problem.jl
+++ /dev/null
@@ -1,242 +0,0 @@
-export UnitaryVariationalProblem
-
-
-"""
-    UnitaryVariationalProblem(
-        system::VariationalQuantumSystem,
-        goal::AbstractPiccoloOperator,
-        T::Int,
-        Δt::Union{Float64, <:AbstractVector{Float64}};
-        robust_times::AbstractVector{<:Union{AbstractVector{Int}, Nothing}}=[nothing for s ∈ system.G_vars],
-        sensitive_times::AbstractVector{<:Union{AbstractVector{Int}, Nothing}}=[nothing for s ∈ system.G_vars],
-        kwargs...
-    )
-
-Constructs a unitary variational problem for optimizing quantum control trajectories.
-
-# Arguments
-
-- `system::VariationalQuantumSystem`: The quantum system to be controlled, containing variational parameters.
-- `goal::AbstractPiccoloOperator`: The target operator or state to achieve at the end of the trajectory.
-- `T::Int`: The total number of timesteps in the trajectory.
-- `Δt::Union{Float64, <:AbstractVector{Float64}}`: The timestep duration or a vector of timestep durations.
-
-# Keyword Arguments
-
-- `robust_times::AbstractVector`: Times at which robustness to variations in the trajectory is enforced.
-- `sensitive_times::AbstractVector`: Times at which sensitivity to variations in the trajectory is enhanced.
-- `unitary_integrator`: The integrator used for unitary evolution (default: `VariationalUnitaryIntegrator`).
-- `state_name::Symbol`: The name of the state variable in the trajectory (default: `:Ũ⃗`).
-- `variational_state_name::Symbol`: The name of the variational state variable (default: `:Ũ⃗ᵥ`).
-- `variational_scales::AbstractVector`: Scaling factors for the variational state variables (default: `1.0`).
-- `control_name::Symbol`: The name of the control variable (default: `:u`).
-- `timestep_name::Symbol`: The name of the timestep variable (default: `:Δt`).
-- `init_trajectory::Union{NamedTrajectory, Nothing}`: An optional initial trajectory to start optimization.
-- `u_bound::Float64`: The bound for the control variable `u` (default: `1.0`).
-- `u_bounds`: Bounds for each control variable (default: filled with `u_bound`).
-- `du_bound::Float64`: The bound for the derivative of the control variable (default: `Inf`).
-- `du_bounds`: Bounds for each derivative of the control variable.
-- `ddu_bound::Float64`: The bound for the second derivative of the control variable (default: `1.0`).
-- `ddu_bounds`: Bounds for each second derivative of the control variable.
-- `Δt_min::Float64`: Minimum allowed timestep duration.
-- `Δt_max::Float64`: Maximum allowed timestep duration.
-- `Q::Float64`: Weight for the unitary infidelity objective (default: `100.0`).
-- `Q_v::Float64`: Weight for sensitivity objectives (default: `1.0`).
-- `R`: Regularization weight for control variables (default: `1e-2`).
-- `R_u`, `R_du`, `R_ddu`: Regularization weights for control, its derivative, and second derivative.
-- `constraints::Vector`: Additional constraints for the optimization problem.
-- `piccolo_options::PiccoloOptions`: Options for configuring the Piccolo optimization framework.
-
-# Returns
-
-A `DirectTrajOptProblem` object representing the optimization problem, including the
-trajectory, objective, integrators, and constraints.
-
-# Notes
-
-This function constructs a trajectory optimization problem for quantum control using
-variational principles. It supports robust and sensitive trajectory design, regularization,
-and optional constraints. The problem is solved using the Piccolo optimization framework.
-
-"""
-function UnitaryVariationalProblem end
-
-function UnitaryVariationalProblem(
-    system::VariationalQuantumSystem,
-    goal::AbstractPiccoloOperator,
-    T::Int,
-    Δt::Union{Float64, <:AbstractVector{Float64}};
-    robust_times::AbstractVector{<:AbstractVector{Int}}=[Int[] for s ∈ system.G_vars],
-    sensitive_times::AbstractVector{<:AbstractVector{Int}}=[Int[] for s ∈ system.G_vars],
-    variational_integrator=VariationalUnitaryIntegrator,
-    variational_scales::AbstractVector{<:Float64}=fill(1.0, length(system.G_vars)),
-    state_name::Symbol = :Ũ⃗,
-    variational_state_name::Symbol = :Ũ⃗ᵥ,
-    control_name::Symbol = :u,
-    timestep_name::Symbol = :Δt,
-    init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
-    u_guess::Union{AbstractMatrix{Float64}, Nothing}=nothing,
-    du_bound::Float64=Inf,
-    du_bounds=fill(du_bound, system.n_drives),
-    ddu_bound::Float64=1.0,
-    ddu_bounds=fill(ddu_bound, system.n_drives),
-    Δt_min::Float64=0.5 * minimum(Δt),
-    Δt_max::Float64=2.0 * maximum(Δt),
-    Q::Float64=100.0,
-    Q_s::Float64=1e-2,
-    Q_r::Float64=100.0,
-    R=1e-2,
-    R_u::Union{Float64, Vector{Float64}}=R,
-    R_du::Union{Float64, Vector{Float64}}=R,
-    R_ddu::Union{Float64, Vector{Float64}}=R,
-    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
-    piccolo_options::PiccoloOptions=PiccoloOptions(),
-)
-    if piccolo_options.verbose
-        println("    constructing UnitaryVariationalProblem...")
-        println("\tusing integrator: $(typeof(variational_integrator))")
-        println("\ttotal variational parameters: $(length(system.G_vars))")
-        if !isempty(robust_times)
-            println("\trobust knot points: $(robust_times)")
-        end
-        if !isempty(sensitive_times)
-            println("\tsensitive knot points: $(sensitive_times)")
-        end
-    end
-
-    # Trajectory
-    if !isnothing(init_trajectory)
-        traj = deepcopy(init_trajectory)
-    else
-        traj = initialize_trajectory(
-            goal,
-            T,
-            Δt,
-            system.n_drives,
-            (system.drive_bounds, du_bounds, ddu_bounds);
-            state_name=state_name,
-            control_name=control_name,
-            timestep_name=timestep_name,
-            Δt_bounds=(Δt_min, Δt_max),
-            zero_initial_and_final_derivative=piccolo_options.zero_initial_and_final_derivative,
-            geodesic=piccolo_options.geodesic,
-            bound_state=piccolo_options.bound_state,
-            u_guess=u_guess,
-            rollout_integrator=piccolo_options.rollout_integrator,
-            verbose=piccolo_options.verbose
-        )
-    end
-
-    _, Ũ⃗_vars =  variational_unitary_rollout(
-        traj, 
-        system;
-        unitary_name=state_name,
-        drive_name=control_name
-    )
-
-    # Add variational components to the trajectory
-    var_state_names = Tuple(
-        Symbol(string(variational_state_name) * "$(i)") for i in eachindex(system.G_vars)
-    )
-    var_comps_inits = NamedTuple{var_state_names}(
-        Ũ⃗_v[:, 1] / scale for (scale, Ũ⃗_v) in zip(variational_scales, Ũ⃗_vars)
-    )
-    var_comps_data = NamedTuple{var_state_names}(
-        Ũ⃗_v / scale for (scale, Ũ⃗_v) in zip(variational_scales, Ũ⃗_vars)
-    )
-    traj = add_components(
-        traj, 
-        var_comps_data; 
-        type=:state,
-        initial=merge(traj.initial, var_comps_inits)
-    )
-
-    control_names = [
-        name for name ∈ traj.names
-            if endswith(string(name), string(control_name))
-    ]
-
-    # objective
-    J = UnitaryInfidelityObjective(goal, state_name, traj; Q=Q)
-    J += QuadraticRegularizer(control_names[1], traj, R_u)
-    J += QuadraticRegularizer(control_names[2], traj, R_du)
-    J += QuadraticRegularizer(control_names[3], traj, R_ddu)
-
-    # sensitivity
-    for (name, scale, s, r) ∈ zip(
-        var_state_names, 
-        variational_scales, 
-        sensitive_times, 
-        robust_times
-    )
-        @assert isdisjoint(s, r)
-        J += UnitarySensitivityObjective(
-            name, 
-            traj, 
-            [s; r]; 
-            Qs=[fill(-Q_s, length(s)); fill(Q_r, length(r))], 
-            scale=scale
-        )
-    end
-
-    # Optional Piccolo constraints and objectives
-    J += apply_piccolo_options!(
-        piccolo_options, constraints, traj;
-        state_names=state_name,
-        state_leakage_indices=goal isa EmbeddedOperator ? 
-            get_iso_vec_leakage_indices(goal) :
-            nothing
-    )
-
-    integrators = [
-        variational_integrator(
-            system, traj, state_name, [var_state_names...], control_name, 
-            scales=variational_scales
-        ),
-        DerivativeIntegrator(traj, control_name, control_names[2]),
-        DerivativeIntegrator(traj, control_names[2], control_names[3]),
-    ]
-
-    return DirectTrajOptProblem(
-        traj,
-        J,
-        integrators;
-        constraints=constraints
-    )
-end
-
-# *************************************************************************** #
-
-@testitem "Sensitive and robust" begin
-    using LinearAlgebra
-    using PiccoloQuantumObjects
-    T_max = 1.0
-    u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-
-    system = QuantumSystem([PAULIS.X, PAULIS.Y], T_max, u_bounds)
-    varsys = VariationalQuantumSystem([PAULIS.X, PAULIS.Y], [PAULIS.X], T_max, u_bounds)
-    T = 50
-    Δt = 0.2
-
-    sense_scale = 8.0
-    sense_prob = UnitaryVariationalProblem(
-        varsys, GATES.X, T, Δt, 
-        variational_scales=[sense_scale], 
-        sensitive_times=[[T]],
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    solve!(sense_prob, max_iter=20, print_level=1, verbose=false)
-
-    rob_scale = 1 / 8.0
-    rob_prob = UnitaryVariationalProblem(
-        varsys, GATES.X, T, Δt, 
-        variational_scales=[rob_scale], 
-        robust_times=[[T]],
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
-    solve!(rob_prob, max_iter=20, print_level=1, verbose=false)
-
-    sense_n = norm(sense_scale * sense_prob.trajectory.Ũ⃗ᵥ1)
-    rob_n = norm(rob_scale * rob_prob.trajectory.Ũ⃗ᵥ1)
-    @test sense_n > rob_n
-end
diff --git a/src/quantum_control_problem.jl b/src/quantum_control_problem.jl
new file mode 100644
index 00000000..3cf69c65
--- /dev/null
+++ b/src/quantum_control_problem.jl
@@ -0,0 +1,135 @@
+module QuantumControlProblems
+
+using DirectTrajOpt
+using NamedTrajectories
+using PiccoloQuantumObjects
+
+import PiccoloQuantumObjects: get_trajectory, get_system, get_goal, get_state_name, get_control_name
+import DirectTrajOpt.Solvers: solve!
+
+export QuantumControlProblem
+export get_trajectory, get_system, get_goal, get_state_name, get_control_name
+export solve!
+# Note: solve! is NOT exported to avoid ambiguity with SciMLBase.solve!
+# Users should use: using DirectTrajOpt (to get solve!)
+
+"""
+    QuantumControlProblem{QT<:AbstractQuantumTrajectory}
+
+Wrapper combining quantum trajectory information with trajectory optimization problem.
+
+This type enables:
+- Type-stable dispatch on quantum trajectory type (Unitary, Ket, Density)
+- Clean separation of quantum information (system, goal) from optimization details
+- Composable problem transformations (e.g., SmoothPulseProblem → MinimumTimeProblem)
+
+# Fields
+- `qtraj::QT`: Quantum trajectory containing system, goal, and quantum state information
+- `prob::DirectTrajOptProblem`: Direct trajectory optimization problem with objective, dynamics, constraints
+
+# Construction
+Typically created via problem templates:
+```julia
+qtraj = UnitaryTrajectory(sys, U_goal, N)
+qcp = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+```
+
+# Accessors
+- `get_trajectory(qcp)`: Get the NamedTrajectory
+- `get_system(qcp)`: Get the QuantumSystem
+- `get_goal(qcp)`: Get the goal state/unitary
+- `get_state_name(qcp)`: Get the state variable name
+- `get_control_name(qcp)`: Get the control variable name
+
+# Solving
+```julia
+solve!(qcp; max_iter=100, verbose=true)
+```
+"""
+struct QuantumControlProblem{QT<:AbstractQuantumTrajectory}
+    qtraj::QT
+    prob::DirectTrajOptProblem
+end
+
+# ============================================================================= #
+# Convenience accessors - extend PiccoloQuantumObjects methods
+# ============================================================================= #
+
+"""
+    get_trajectory(qcp::QuantumControlProblem)
+
+Get the NamedTrajectory from the optimization problem.
+"""
+get_trajectory(qcp::QuantumControlProblem) = qcp.prob.trajectory
+
+"""
+    get_system(qcp::QuantumControlProblem)
+
+Get the QuantumSystem from the quantum trajectory.
+"""
+get_system(qcp::QuantumControlProblem) = get_system(qcp.qtraj)
+
+"""
+    get_goal(qcp::QuantumControlProblem)
+
+Get the goal state/operator from the quantum trajectory.
+"""
+get_goal(qcp::QuantumControlProblem) = get_goal(qcp.qtraj)
+
+"""
+    get_state_name(qcp::QuantumControlProblem)
+
+Get the state variable name from the quantum trajectory.
+"""
+get_state_name(qcp::QuantumControlProblem) = get_state_name(qcp.qtraj)
+
+"""
+    get_control_name(qcp::QuantumControlProblem)
+
+Get the control variable name from the quantum trajectory.
+"""
+get_control_name(qcp::QuantumControlProblem) = get_control_name(qcp.qtraj)
+
+# ============================================================================= #
+# Forward DirectTrajOptProblem methods
+# ============================================================================= #
+
+"""
+    solve!(qcp::QuantumControlProblem; kwargs...)
+
+Solve the quantum control problem by forwarding to the inner DirectTrajOptProblem.
+
+All keyword arguments are passed to the DirectTrajOpt solver.
+"""
+solve!(qcp::QuantumControlProblem; kwargs...) = solve!(qcp.prob; kwargs...)
+# TODO: make sure qtraj is updated after solve!
+
+# Forward other common DirectTrajOptProblem accessors
+Base.getproperty(qcp::QuantumControlProblem, s::Symbol) = begin
+    if s === :qtraj
+        getfield(qcp, :qtraj)
+    elseif s === :prob
+        getfield(qcp, :prob)
+    # Forward to prob for common fields
+    elseif s in (:objective, :dynamics, :constraints, :trajectory)
+        getproperty(qcp.prob, s)
+    else
+        # Fall back to default behavior
+        getfield(qcp, s)
+    end
+end
+
+# ============================================================================= #
+# Display
+# ============================================================================= #
+
+function Base.show(io::IO, qcp::QuantumControlProblem{QT}) where {QT}
+    println(io, "QuantumControlProblem{$QT}")
+    println(io, "  System: $(typeof(get_system(qcp)))")
+    println(io, "  Goal: $(typeof(get_goal(qcp)))")
+    println(io, "  Trajectory: $(qcp.prob.trajectory.N) knots")
+    println(io, "  State: $(get_state_name(qcp))")
+    print(io, "  Controls: $(get_control_name(qcp))")
+end
+
+end
diff --git a/src/quantum_integrators.jl b/src/quantum_integrators.jl
index c711ed90..10b288a5 100644
--- a/src/quantum_integrators.jl
+++ b/src/quantum_integrators.jl
@@ -1,15 +1,18 @@
 module QuantumIntegrators
 
-export KetIntegrator
-export UnitaryIntegrator
-export DensityMatrixIntegrator
-export VariationalUnitaryIntegrator
-
 using LinearAlgebra
 using NamedTrajectories
 using DirectTrajOpt
 using PiccoloQuantumObjects
+using PiccoloQuantumObjects: SamplingTrajectory, EnsembleTrajectory, 
+    get_ensemble_state_names, build_sampling_trajectory, build_ensemble_trajectory_from_trajectories
 using SparseArrays
+using TestItems
+
+import DirectTrajOpt: BilinearIntegrator
+
+# Import QuantumTrajectories types (will be loaded before this module)
+using ..QuantumTrajectories
 
 const ⊗ = kron
 
@@ -17,33 +20,149 @@ const ⊗ = kron
 # Default Integrators
 # ----------------------------------------------------------------------------- #
 
-function KetIntegrator(
-    sys::QuantumSystem,
+# Dispatch on quantum trajectory types
+function BilinearIntegrator(qtraj::UnitaryTrajectory)
+    sys = get_system(qtraj)
+    traj = get_trajectory(qtraj)
+    Ĝ = u_ -> I(sys.levels) ⊗ sys.G(u_, 0.0)
+    return BilinearIntegrator(Ĝ, get_state_name(qtraj), get_control_name(qtraj), traj)
+end
+
+function BilinearIntegrator(qtraj::KetTrajectory)
+    sys = get_system(qtraj)
+    traj = get_trajectory(qtraj)
+    Ĝ = u_ -> sys.G(u_, 0.0)
+    return BilinearIntegrator(Ĝ, get_state_name(qtraj), get_control_name(qtraj), traj)
+end
+
+function BilinearIntegrator(qtraj::DensityTrajectory)
+    sys = get_system(qtraj)
+    traj = get_trajectory(qtraj)
+    return BilinearIntegrator(sys.𝒢, get_state_name(qtraj), get_control_name(qtraj), traj)
+end
+
+# ----------------------------------------------------------------------------- #
+# SamplingTrajectory Integrators
+# ----------------------------------------------------------------------------- #
+
+"""
+    BilinearIntegrator(qtraj::SamplingTrajectory, traj::NamedTrajectory)
+
+Create a vector of BilinearIntegrators for each system in a SamplingTrajectory.
+
+Each system in the sampling ensemble gets its own dynamics integrator, but they
+all share the same control variables. The trajectory `traj` should be the 
+expanded trajectory with sample state variables (e.g., from `build_sampling_trajectory`).
+
+# Returns
+- `Vector{BilinearIntegrator}`: One integrator per system in the ensemble
+"""
+function BilinearIntegrator(qtraj::SamplingTrajectory, traj::NamedTrajectory)
+    state_names = get_ensemble_state_names(qtraj)
+    control_sym = get_control_name(qtraj)
+    systems = qtraj.systems
+    
+    return [
+        _sampling_integrator(qtraj.base_trajectory, sys, traj, name, control_sym)
+        for (sys, name) in zip(systems, state_names)
+    ]
+end
+
+# Helper to create single integrator for sampling - dispatches on base trajectory type
+function _sampling_integrator(
+    base_qtraj::UnitaryTrajectory,
+    sys::AbstractQuantumSystem,
     traj::NamedTrajectory,
-    ψ̃::Symbol,
-    u::Symbol
+    state_sym::Symbol,
+    control_sym::Symbol
+)
+    Ĝ = u_ -> I(sys.levels) ⊗ sys.G(u_, 0.0)
+    return BilinearIntegrator(Ĝ, state_sym, control_sym, traj)
+end
+
+function _sampling_integrator(
+    base_qtraj::KetTrajectory,
+    sys::AbstractQuantumSystem,
+    traj::NamedTrajectory,
+    state_sym::Symbol,
+    control_sym::Symbol
 )
     Ĝ = u_ -> sys.G(u_, 0.0)
-    return BilinearIntegrator(Ĝ, traj, ψ̃, u)
+    return BilinearIntegrator(Ĝ, state_sym, control_sym, traj)
+end
+
+function _sampling_integrator(
+    base_qtraj::DensityTrajectory,
+    sys::OpenQuantumSystem,
+    traj::NamedTrajectory,
+    state_sym::Symbol,
+    control_sym::Symbol
+)
+    return BilinearIntegrator(sys.𝒢, state_sym, control_sym, traj)
 end
 
-function UnitaryIntegrator(
-    sys::QuantumSystem,
+# ----------------------------------------------------------------------------- #
+# EnsembleTrajectory Integrators
+# ----------------------------------------------------------------------------- #
+
+"""
+    BilinearIntegrator(qtraj::EnsembleTrajectory, traj::NamedTrajectory)
+
+Create a vector of BilinearIntegrators for each trajectory in an EnsembleTrajectory.
+
+Each trajectory in the ensemble gets its own dynamics integrator evolving under
+the shared system, but with different state variables. The trajectory `traj` should
+be the expanded trajectory with ensemble state variables (e.g., from 
+`build_ensemble_trajectory_from_trajectories`).
+
+# Returns
+- `Vector{BilinearIntegrator}`: One integrator per trajectory in the ensemble
+"""
+function BilinearIntegrator(qtraj::EnsembleTrajectory, traj::NamedTrajectory)
+    state_names = get_ensemble_state_names(qtraj)
+    control_sym = get_control_name(qtraj)
+    sys = get_system(qtraj)  # Shared system for all
+    
+    # Use the base trajectory type to determine integrator construction
+    base_type = eltype(qtraj.trajectories)
+    
+    return [
+        _ensemble_integrator(base_type, sys, traj, name, control_sym)
+        for name in state_names
+    ]
+end
+
+# Helper to create single integrator for ensemble - dispatches on trajectory type
+function _ensemble_integrator(
+    ::Type{<:UnitaryTrajectory},
+    sys::AbstractQuantumSystem,
     traj::NamedTrajectory,
-    Ũ⃗::Symbol,
-    u::Symbol
+    state_sym::Symbol,
+    control_sym::Symbol
 )
-    Ĝ = u_ -> I(sys.levels) ⊗ sys.G(u_, 0.0)
-    return BilinearIntegrator(Ĝ, traj, Ũ⃗, u)
+    Ĝ = u_ -> I(sys.levels) ⊗ sys.G(u_, 0.0)
+    return BilinearIntegrator(Ĝ, state_sym, control_sym, traj)
 end
 
-function DensityMatrixIntegrator(
+function _ensemble_integrator(
+    ::Type{<:KetTrajectory},
+    sys::AbstractQuantumSystem,
+    traj::NamedTrajectory,
+    state_sym::Symbol,
+    control_sym::Symbol
+)
+    Ĝ = u_ -> sys.G(u_, 0.0)
+    return BilinearIntegrator(Ĝ, state_sym, control_sym, traj)
+end
+
+function _ensemble_integrator(
+    ::Type{<:DensityTrajectory},
     sys::OpenQuantumSystem,
     traj::NamedTrajectory,
-    ρ̃::Symbol,
-    u::Symbol
+    state_sym::Symbol,
+    control_sym::Symbol
 )
-    return BilinearIntegrator(sys.𝒢, traj, ρ̃, u)
+    return BilinearIntegrator(sys.𝒢, state_sym, control_sym, traj)
 end
 
 # ----------------------------------------------------------------------------- #
@@ -59,26 +178,191 @@ function VariationalKetIntegrator(
     scale::Float64=1.0,
 )
     var_ψ̃ = vcat(ψ̃, ψ̃_variations...)
-    G = a -> Isomorphisms.var_G(sys.G(a), [G(a) / scale for G in sys.G_vars])
-    return BilinearIntegrator(G, traj, var_ψ̃, u)
+    G = u_ -> Isomorphisms.var_G(sys.G(u_), [G(u_) / scale for G in sys.G_vars])
+    return BilinearIntegrator(G, var_ψ̃, u, traj)
 end
 
 function VariationalUnitaryIntegrator(
     sys::VariationalQuantumSystem,
     traj::NamedTrajectory,
-    Ũ⃗::Symbol,
-    Ũ⃗_variations::AbstractVector{Symbol},
+    Ũ⃗::Symbol,
+    Ũ⃗_variations::AbstractVector{Symbol},
     u::Symbol;
     scales::AbstractVector{<:Float64}=fill(1.0, length(sys.G_vars)),
 )
-    var_Ũ⃗ = vcat(Ũ⃗, Ũ⃗_variations...)
+    var_Ũ⃗ = vcat(Ũ⃗, Ũ⃗_variations...)
 
-    function Ĝ(u)
-        G0 = sys.G(u)
-        Gs = typeof(G0)[I(sys.levels) ⊗ G(u) / scale for (scale, G) in zip(scales, sys.G_vars)]
+    function Ĝ(u_)
+        G0 = sys.G(u_)
+        Gs = typeof(G0)[I(sys.levels) ⊗ G(u_) / scale for (scale, G) in zip(scales, sys.G_vars)]
         return Isomorphisms.var_G(I(sys.levels) ⊗ G0, Gs)
     end
-    return BilinearIntegrator(Ĝ, traj, var_Ũ⃗, u)
+    return BilinearIntegrator(Ĝ, var_Ũ⃗, u, traj)
+end
+
+# ----------------------------------------------------------------------------- #
+# Tests
+# ----------------------------------------------------------------------------- #
+
+@testitem "BilinearIntegrator dispatch on UnitaryTrajectory" begin
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    qtraj = UnitaryTrajectory(sys, GATES[:H], 10)
+
+    integrator = BilinearIntegrator(qtraj)
+
+    @test integrator isa BilinearIntegrator
+
+    test_integrator(integrator, get_trajectory(qtraj); atol=1e-3)
+end
+
+@testitem "BilinearIntegrator dispatch on KetTrajectory" begin
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+    qtraj = KetTrajectory(sys, ψ_init, ψ_goal, 10)
+
+    integrator = BilinearIntegrator(qtraj)
+
+    @test integrator isa BilinearIntegrator
+    test_integrator(integrator, get_trajectory(qtraj); atol=1e-3)
+end
+
+@testitem "BilinearIntegrator dispatch on DensityTrajectory" tags=[:experimental] begin
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    sys = OpenQuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 20.0, [1.0, 1.0])
+    ρ_init = ComplexF64[1.0 0.0; 0.0 0.0]
+    ρ_goal = ComplexF64[0.0 0.0; 0.0 1.0]
+    qtraj = DensityTrajectory(sys, ρ_init, ρ_goal, 10)
+
+    integrator = BilinearIntegrator(qtraj)
+
+    @test integrator isa BilinearIntegrator
+    test_integrator(integrator, get_trajectory(qtraj); atol=1e-2, show_hessian_diff=true)
+end
+
+@testitem "BilinearIntegrator dispatch on SamplingTrajectory (Unitary)" begin
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Create systems with parameter variation
+    sys1 = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    sys2 = QuantumSystem(1.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+
+    # Create base trajectory and sampling trajectory
+    base_qtraj = UnitaryTrajectory(sys1, GATES[:H], 10)
+    sampling_qtraj = SamplingTrajectory(base_qtraj, [sys1, sys2])
+
+    # Build expanded trajectory
+    base_traj = get_trajectory(base_qtraj)
+    expanded_traj, state_names = build_sampling_trajectory(base_traj, :Ũ⃗, 2)
+
+    # Create integrators
+    integrators = BilinearIntegrator(sampling_qtraj, expanded_traj)
+
+    @test integrators isa Vector{<:BilinearIntegrator}
+    @test length(integrators) == 2
+
+    # Test each integrator
+    for integrator in integrators
+        test_integrator(integrator, expanded_traj; atol=1e-3)
+    end
+end
+
+@testitem "BilinearIntegrator dispatch on SamplingTrajectory (Ket)" begin
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Create systems with parameter variation
+    sys1 = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+    sys2 = QuantumSystem(1.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+
+    ψ_init = ComplexF64[1.0, 0.0]
+    ψ_goal = ComplexF64[0.0, 1.0]
+
+    # Create base trajectory and sampling trajectory
+    base_qtraj = KetTrajectory(sys1, ψ_init, ψ_goal, 10)
+    sampling_qtraj = SamplingTrajectory(base_qtraj, [sys1, sys2])
+
+    # Build expanded trajectory
+    base_traj = get_trajectory(base_qtraj)
+    expanded_traj, state_names = build_sampling_trajectory(base_traj, :ψ̃, 2)
+
+    # Create integrators
+    integrators = BilinearIntegrator(sampling_qtraj, expanded_traj)
+
+    @test integrators isa Vector{<:BilinearIntegrator}
+    @test length(integrators) == 2
+
+    for integrator in integrators
+        test_integrator(integrator, expanded_traj; atol=1e-3)
+    end
+end
+
+@testitem "BilinearIntegrator dispatch on EnsembleTrajectory (Ket)" begin
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Shared system
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+
+    # Different initial/goal states
+    ψ0 = ComplexF64[1.0, 0.0]
+    ψ1 = ComplexF64[0.0, 1.0]
+
+    qtraj1 = KetTrajectory(sys, ψ0, ψ1, 10)  # |0⟩ → |1⟩
+    qtraj2 = KetTrajectory(sys, ψ1, ψ0, 10)  # |1⟩ → |0⟩
+
+    # Create ensemble trajectory
+    ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+
+    # Build expanded trajectory - pass quantum trajectories, not named trajectories
+    expanded_traj, state_names = build_ensemble_trajectory_from_trajectories([qtraj1, qtraj2])
+
+    # Create integrators
+    integrators = BilinearIntegrator(ensemble_qtraj, expanded_traj)
+
+    @test integrators isa Vector{<:BilinearIntegrator}
+    @test length(integrators) == 2
+
+    for integrator in integrators
+        test_integrator(integrator, expanded_traj; atol=1e-3)
+    end
+end
+
+@testitem "BilinearIntegrator dispatch on EnsembleTrajectory (Unitary)" begin
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+
+    # Shared system
+    sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+
+    # Different target gates
+    qtraj1 = UnitaryTrajectory(sys, GATES[:X], 10)
+    qtraj2 = UnitaryTrajectory(sys, GATES[:H], 10)
+
+    # Create ensemble trajectory
+    ensemble_qtraj = EnsembleTrajectory([qtraj1, qtraj2])
+
+    # Build expanded trajectory - pass quantum trajectories, not named trajectories
+    expanded_traj, state_names = build_ensemble_trajectory_from_trajectories([qtraj1, qtraj2])
+
+    # Create integrators
+    integrators = BilinearIntegrator(ensemble_qtraj, expanded_traj)
+
+    @test integrators isa Vector{<:BilinearIntegrator}
+    @test length(integrators) == 2
+
+    for integrator in integrators
+        test_integrator(integrator, expanded_traj; atol=1e-3)
+    end
 end
 
 
diff --git a/src/quantum_objectives.jl b/src/quantum_objectives.jl
index 5e298aaa..eed8e8e3 100644
--- a/src/quantum_objectives.jl
+++ b/src/quantum_objectives.jl
@@ -1,6 +1,7 @@
 module QuantumObjectives
 
 export KetInfidelityObjective
+export KetsCoherentInfidelityObjective
 export UnitaryInfidelityObjective
 export DensityMatrixPureStateInfidelityObjective
 export UnitarySensitivityObjective
@@ -25,6 +26,11 @@ function ket_fidelity_loss(
     return abs2(ψ_goal' * ψ)
 end 
 
+"""
+    KetInfidelityObjective(ψ̃_name, traj; Q=100.0)
+
+Create a terminal objective for ket state infidelity, using the goal from `traj.goal[ψ̃_name]`.
+"""
 function KetInfidelityObjective(
     ψ̃_name::Symbol,
     traj::NamedTrajectory;
@@ -35,6 +41,70 @@ function KetInfidelityObjective(
     return TerminalObjective(ℓ, ψ̃_name, traj; Q=Q)
 end
 
+"""
+    KetInfidelityObjective(ψ_goal, ψ̃_name, traj; Q=100.0)
+
+Create a terminal objective for ket state infidelity with an explicit goal state.
+
+This variant is useful for SamplingProblem and EnsembleTrajectory where the goal
+is shared across multiple state variables that don't have individual goals in `traj.goal`.
+
+# Arguments
+- `ψ_goal::AbstractVector{<:Complex}`: The target ket state (complex vector)
+- `ψ̃_name::Symbol`: Name of the isomorphic state variable in the trajectory
+- `traj::NamedTrajectory`: The trajectory
+
+# Keyword Arguments
+- `Q::Float64=100.0`: Weight on the infidelity objective
+"""
+function KetInfidelityObjective(
+    ψ_goal::AbstractVector{<:Complex},
+    ψ̃_name::Symbol,
+    traj::NamedTrajectory;
+    Q=100.0
+)
+    ℓ = ψ̃ -> abs(1 - ket_fidelity_loss(ψ̃, ComplexF64.(ψ_goal)))
+    return TerminalObjective(ℓ, ψ̃_name, traj; Q=Q)
+end
+
+function _partitions(dims::AbstractVector{Int})
+    ends = cumsum(dims)
+    return UnitRange.(ends .- dims .+ 1, ends)
+end
+
+"""
+    KetsCoherentInfidelityObjective(ψ_goals, ψ̃_names, traj; Q=100.0)
+
+Phase-sensitive ensemble fidelity objective for a goal unitary via state preparations.
+
+For an ensemble of N prepared states with final states |ϕᵢ⟩ and goals |gᵢ⟩, computes the 
+coherent overlap ℱ = |∑ᵢ ⟨ϕᵢ | gᵢ⟩|² / N². If the prepared states form a complete  basis,
+ℱ reduces to |Tr(U† U_target)|² / N².
+
+# Arguments
+- `ψ_goals::AbstractVector{<:AbstractVector{<:Complex}}`: The target ket states (complex vector)
+- `ψ̃_names::Symbol`: Names of the isomorphic state variables in the trajectory
+- `traj::NamedTrajectory`: The trajectory
+
+# Keyword Arguments
+- `Q::Float64=100.0`: Weight on the infidelity objective
+"""
+function KetsCoherentInfidelityObjective(
+    ψ_goals::AbstractVector{<:AbstractVector{<:Complex}},
+    ψ̃_names::AbstractVector{Symbol}, 
+    traj::NamedTrajectory;
+    Q=100.0
+)
+    N = length(ψ̃_names)
+    @assert length(ψ_goals) == N
+    idxs = _partitions([traj.dims[n] for n in ψ̃_names])
+    function ℓ(x)
+        ℱ = abs2(sum(iso_to_ket(x[idx])'g for (idx, g) in zip(idxs, ψ_goals))) / N^2
+        return abs(1 - ℱ)
+    end
+    return KnotPointObjective(ℓ, ψ̃_names, traj, times=[traj.N], Qs=[Q])
+end
+
 
 # ---------------------------------------------------------
 #                       Unitaries
@@ -176,5 +246,10 @@ function LeakageObjective(
     )
 end
 
+# ---------------------------------------------------------
+#                       Tests
+# ---------------------------------------------------------
+
+using TestItems
 
 end
\ No newline at end of file
diff --git a/src/quantum_system_templates/_quantum_system_templates.jl b/src/quantum_system_templates/_quantum_system_templates.jl
deleted file mode 100644
index c820b3a5..00000000
--- a/src/quantum_system_templates/_quantum_system_templates.jl
+++ /dev/null
@@ -1,13 +0,0 @@
-module QuantumSystemTemplates
-
-using PiccoloQuantumObjects
-using LinearAlgebra
-using TestItems
-
-const ⊗ = kron
-
-include("transmons.jl")
-include("rydberg.jl")
-include("cats.jl")
-
-end
diff --git a/src/quantum_system_templates/cats.jl b/src/quantum_system_templates/cats.jl
deleted file mode 100644
index e7d01d7a..00000000
--- a/src/quantum_system_templates/cats.jl
+++ /dev/null
@@ -1,56 +0,0 @@
-export CatSystem
-export coherent_ket
-export get_cat_controls
-
-function coherent_ket(α::Union{Real, Complex}, levels::Int)::Vector{ComplexF64}
-    return [exp(-0.5 * abs2(α)) * α^n / sqrt(factorial(n)) for n in 0:levels-1]
-end
-
-"""
-    CatSystem(;
-        g2::Real=0.36,
-        χ_aa::Real=-7e-3,
-        χ_bb::Real=-32,
-        χ_ab::Real=0.79,
-        κa::Real=53e-3,
-        κb::Real=13,
-        cat_levels::Int=13,
-        buffer_levels::Int=3,
-        prefactor::Real=1,
-    )::OpenQuantumSystem
-
-Returns an `OpenQuantumSystem` for a quantum cat.
-"""
-function CatSystem(;
-    g2::Real=0.36,
-    χ_aa::Real=-7e-3,
-    χ_bb::Real=-32,
-    χ_ab::Real=0.79,
-    κa::Real=53e-3,
-    κb::Real=13,
-    cat_levels::Int=13,
-    buffer_levels::Int=3,
-    prefactor::Real=1,
-)
-
-    # Cat ⊗ Buffer
-    a = annihilate(cat_levels) ⊗ Matrix(1.0I, buffer_levels, buffer_levels)
-    b = Matrix(1.0I, cat_levels, cat_levels) ⊗ annihilate(buffer_levels)
-
-    H_drift = -χ_aa/2 * a'a'a*a - χ_bb/2 * b'b'b*b - χ_ab * a'a*b'b + g2 * a'a'b + conj(g2) * a*a*b'
-
-    # buffer drive, kerr-correction drive
-    H_drives = [b + b', a'a]
-
-    L_dissipators = [√κa * a, √κb * b]
-
-    H_drift *= 2π
-    H_drives .*= 2π
-    L_dissipators .*= sqrt(2π)
-
-    return OpenQuantumSystem(
-        H_drift,
-        H_drives,
-        L_dissipators;
-    )
-end
diff --git a/src/quantum_system_templates/rydberg.jl b/src/quantum_system_templates/rydberg.jl
deleted file mode 100644
index 0b07ee5d..00000000
--- a/src/quantum_system_templates/rydberg.jl
+++ /dev/null
@@ -1,146 +0,0 @@
-export RydbergChainSystem
-
-function generate_pattern(N::Int, i::Int)
-    # Create an array filled with 'I'
-    qubits = fill('I', N)
-    # Insert 'n' at position i and i+1, ensuring it doesn't exceed the array bounds
-    if i <= N && i+1 <= N
-        qubits[i] = 'n'
-        qubits[i+1] = 'n'
-    end
-    return join(qubits)
-end
-
-function generate_pattern_with_gap(N::Int, i::Int, gap::Int)
-    # Create an array filled with 'I'
-    qubits = fill('I', N)
-    # Insert 'n' at position i and i+gap+1, ensuring it doesn't exceed the array bounds
-    if i <= N && i+gap+1 <= N
-        qubits[i] = 'n'
-        qubits[i+gap+1] = 'n'
-    end
-    return join(qubits)
-end
-
-"""
-    lift(x::Char, i::Int, N::Int)::String
-
-Embed a character into a string of the form 'I' * N at a specific position (meant for use with `PiccoloQuantumObjects.QuantumObjectUtils.operator_from_string`).
-"""
-function lift(x::Char,i::Int, N::Int)
-    qubits = fill('I', N)
-    qubits[i] = x
-    return join(qubits)
-end
-
-@doc raw"""
-    RydbergChainSystem(;
-        N::Int=3, # number of atoms
-        C::Float64=862690*2π,
-        distance::Float64=10.0, # μm
-        cutoff_order::Int=2, # 1 is nearest neighbor, 2 is next-nearest neighbor, etc.
-        local_detune::Bool=false, # If true, include one local detuning pattern.
-        all2all::Bool=true, # If true, include all-to-all interactions.
-        ignore_Y_drive::Bool=false, # If true, ignore the Y drive. (In the experiments, X&Y drives are implemented by Rabi amplitude and its phase.)
-    )::QuantumSystem
-
-Returns a `QuantumSystem` object for the Rydberg atom chain in the spin basis
-    |g⟩ = |0⟩ = [1, 0], |r⟩ = |1⟩ = [0, 1].
-
-```math
-H = \sum_i 0.5*\Omega_i(t)\cos(\phi_i(t)) \sigma_i^x - 0.5*\Omega_i(t)\sin(\phi_i(t)) \sigma_i^y - \sum_i \Delta_i(t)n_i + \sum_{i QuantumSystem
-
-Returns a `QuantumSystem` object for a transmon qubit, with the Hamiltonian
-
-```math
-H = \omega a^\dagger a - \frac{\delta}{2} a^\dagger a^\dagger a a
-```
-
-where `a` is the annihilation operator.
-
-# Keyword Arguments
-- `ω`: The frequency of the transmon, in GHz.
-- `δ`: The anharmonicity of the transmon, in GHz.
-- `levels`: The number of levels in the transmon.
-- `lab_frame`: Whether to use the lab frame Hamiltonian, or an ω-rotating frame.
-- `frame_ω`: The frequency of the rotating frame, in GHz.
-- `mutiply_by_2π`: Whether to multiply the Hamiltonian by 2π, set to true by default because the frequency is in GHz.
-- `lab_frame_type`: The type of lab frame Hamiltonian to use, one of (:duffing, :quartic, :cosine).
-- `drives`: Whether to include drives in the Hamiltonian.
-"""
-function TransmonSystem(;
-    T_max::Float64=10.0,
-    drive_bounds::Vector{<:Union{Tuple{Float64, Float64}, Float64}}=fill(1.0, 2),
-    ω::Float64=4.0,  # GHz
-    δ::Float64=0.2, # GHz
-    levels::Int=3,
-    lab_frame::Bool=false,
-    frame_ω::Float64=lab_frame ? 0.0 : ω,
-    mutiply_by_2π::Bool=true,
-    lab_frame_type::Symbol=:duffing,
-    drives::Bool=true,
-    kwargs...
-)
-
-    @assert lab_frame_type ∈ (:duffing, :quartic, :cosine) "lab_frame_type must be one of (:duffing, :quartic, :cosine)"
-
-    if lab_frame
-        if frame_ω ≉ 0.0
-            frame_ω = 0.0
-        end
-    end
-
-    if frame_ω ≉ 0.0
-        lab_frame = false
-    end
-
-    a = annihilate(levels)
-
-    if lab_frame
-        if lab_frame_type == :duffing
-            H_drift = ω * a' * a - δ / 2 * a' * a' * a * a
-        elseif lab_frame_type == :quartic
-            ω₀ = ω + δ
-            H_drift = ω₀ * a' * a - δ / 12 * (a + a')^4
-        elseif lab_frame_type == :cosine
-            ω₀ = ω + δ
-            E_C = δ
-            E_J = ω₀^2 / 8E_C
-            n̂ = im / 2 * (E_J / 2E_C)^(1/4) * (a - a')
-            φ̂ = (2E_C / E_J)^(1/4) * (a + a')
-            H_drift = 4 * E_C * n̂^2 - E_J * cos(φ̂)
-            # H_drift = 4 * E_C * n̂^2 - E_J * (I - φ̂^2 / 2 + φ̂^4 / 24)
-        end
-    else
-        H_drift = (ω - frame_ω) * a' * a - δ / 2 * a' * a' * a * a
-    end
-
-    if drives
-        H_drives = [a + a', 1.0im * (a - a')]
-    else
-        H_drives = Matrix{ComplexF64}[]
-    end
-
-    if mutiply_by_2π
-        H_drift *= 2π
-        H_drives .*= 2π
-    end
-
-    return QuantumSystem(
-        H_drift,
-        H_drives,
-        T_max,
-        drive_bounds;
-    )
-end
-
-struct QuantumSystemCoupling
-    op
-    g_ij
-    pair
-    subsystem_levels
-    coupling_type
-end
-
-@doc raw"""
-    TransmonDipoleCoupling(
-        g_ij::Float64,
-        pair::Tuple{Int, Int},
-        subsystem_levels::Vector{Int};
-        lab_frame::Bool=false,
-    ) -> QuantumSystemCoupling
-
-    TransmonDipoleCoupling(
-        g_ij::Float64,
-        pair::Tuple{Int, Int},
-        sub_systems::Vector{QuantumSystem};
-        kwargs...
-    ) -> QuantumSystemCoupling
-
-Returns a `QuantumSystemCoupling` object for a transmon qubit. In the lab frame, the Hamiltonian coupling term is
-
-```math
-H = g_{ij} (a_i + a_i^\dagger) (a_j + a_j^\dagger)
-```
-
-In the rotating frame, the Hamiltonian coupling term is
-
-```math
-H = g_{ij} (a_i a_j^\dagger + a_i^\dagger a_j)
-```
-
-where `a_i` is the annihilation operator for the `i`th transmon.
-
-"""
-function TransmonDipoleCoupling end
-
-function TransmonDipoleCoupling(
-    g_ij::Float64,
-    pair::Tuple{Int, Int},
-    subsystem_levels::Vector{Int};
-    lab_frame::Bool=false,
-    mulitply_by_2π::Bool=true,
-)
-
-    i, j = pair
-    a_i = lift_operator(annihilate(subsystem_levels[i]), i, subsystem_levels)
-    a_j = lift_operator(annihilate(subsystem_levels[j]), j, subsystem_levels)
-
-    if lab_frame
-        op = g_ij * (a_i + a_i') * (a_j + a_j')
-    else
-        op = g_ij * (a_i * a_j' + a_i' * a_j)
-    end
-
-    if mulitply_by_2π
-        op *= 2π
-    end
-
-    return QuantumSystemCoupling(
-        op,
-        g_ij,
-        pair,
-        subsystem_levels,
-        TransmonDipoleCoupling,
-    )
-end
-
-function TransmonDipoleCoupling(
-    g_ij::Float64,
-    pair::Tuple{Int, Int},
-    sub_systems::AbstractVector{<:AbstractQuantumSystem};
-    kwargs...
-)
-    subsystem_levels = [sys.levels for sys ∈ sub_systems]
-    return TransmonDipoleCoupling(g_ij, pair, subsystem_levels; kwargs...)
-end
-
-"""
-    MultiTransmonSystem(
-        ωs::AbstractVector{Float64},
-        δs::AbstractVector{Float64},
-        gs::AbstractMatrix{Float64};
-        levels_per_transmon::Int = 3,
-        subsystem_levels::AbstractVector{Int} = fill(levels_per_transmon, length(ωs)),
-        lab_frame=false,
-        subsystems::AbstractVector{Int} = 1:length(ωs),
-        subsystem_drive_indices::AbstractVector{Int} = 1:length(ωs),
-        kwargs...
-    ) -> CompositeQuantumSystem
-
-Returns a `CompositeQuantumSystem` object for a multi-transmon system.
-"""
-function MultiTransmonSystem(
-    ωs::AbstractVector{Float64},
-    δs::AbstractVector{Float64},
-    gs::AbstractMatrix{Float64};
-    T_max::Float64=10.0,
-    drive_bounds::Union{Float64, Vector{<:Union{Tuple{Float64, Float64}, Float64}}}=[(-1.0, 1.0) for _ in 1:length(ωs)],
-    levels_per_transmon::Int = 3,
-    subsystem_levels::AbstractVector{Int} = fill(levels_per_transmon, length(ωs)),
-    lab_frame=false,
-    subsystems::AbstractVector{Int} = 1:length(ωs),
-    subsystem_drive_indices::AbstractVector{Int} = 1:length(ωs),
-    kwargs...
-)
-    n_subsystems = length(ωs)
-    @assert length(δs) == n_subsystems
-    @assert size(gs) == (n_subsystems, n_subsystems)
-
-    systems = QuantumSystem[]
-
-    for (i, (ω, δ, levels)) ∈ enumerate(zip(ωs, δs, subsystem_levels))
-        if i ∈ subsystems
-            sysᵢ = TransmonSystem(
-                levels=levels,
-                ω=ω,
-                δ=δ,
-                lab_frame=lab_frame,
-                drives=i ∈ subsystem_drive_indices
-            )
-            push!(systems, sysᵢ)
-        end
-    end
-
-    couplings = QuantumSystemCoupling[]
-
-    for local_i = 1:length(systems)-1
-        for local_j = local_i+1:length(systems)
-            global_i = subsystems[local_i]
-            global_j = subsystems[local_j]
-            push!(
-                couplings,
-                TransmonDipoleCoupling(gs[global_i, global_j], (local_i, local_j), [sys.levels for sys in systems]; lab_frame=lab_frame)
-            )
-        end
-    end
-
-    levels = prod([sys.levels for sys in systems])
-    H_drift = sum(c -> c.op, couplings; init=zeros(ComplexF64, levels, levels))
-    return CompositeQuantumSystem(H_drift, systems, T_max, drive_bounds)
-end
-
-# *************************************************************************** #
-
-@testitem "TransmonSystem: default and custom parameters" begin
-    using PiccoloQuantumObjects
-    sys = TransmonSystem()
-    @test sys isa QuantumSystem
-    @test sys.levels == 3
-    @test sys.n_drives == 2
-
-    sys2 = TransmonSystem(ω=5.0, δ=0.3, levels=4, lab_frame=true, frame_ω=0.0, lab_frame_type=:duffing, drives=false)
-    @test sys2.levels == 4
-    @test sys2.n_drives == 0
-end
-
-@testitem "TransmonSystem: lab_frame_type variations" begin
-    using PiccoloQuantumObjects
-    sys_duffing = TransmonSystem(lab_frame=true, lab_frame_type=:duffing)
-    sys_quartic = TransmonSystem(lab_frame=true, lab_frame_type=:quartic)
-    sys_cosine = TransmonSystem(lab_frame=true, lab_frame_type=:cosine)
-    @test sys_duffing isa QuantumSystem
-    @test sys_quartic isa QuantumSystem
-    @test sys_cosine isa QuantumSystem
-end
-
-@testitem "TransmonSystem: error on invalid lab_frame_type" begin
-    @test_throws AssertionError TransmonSystem(lab_frame=true, lab_frame_type=:invalid)
-end
-
-@testitem "TransmonDipoleCoupling: both constructors and frames" begin
-    levels = [3, 3]
-    g = 0.01
-
-    c1 = TransmonDipoleCoupling(g, (1, 2), levels, lab_frame=false)
-    c2 = TransmonDipoleCoupling(g, (1, 2), levels, lab_frame=true)
-    @test c1 isa QuantumSystemCoupling
-    @test c2 isa QuantumSystemCoupling
-
-    sys1 = TransmonSystem(levels=3)
-    sys2 = TransmonSystem(levels=3)
-    c3 = TransmonDipoleCoupling(g, (1, 2), [sys1, sys2], lab_frame=false)
-    @test c3 isa QuantumSystemCoupling
-end
-
-@testitem "MultiTransmonSystem: minimal and custom" begin
-    using LinearAlgebra: norm
-    using PiccoloQuantumObjects
-
-    ωs = [4.0, 4.1]
-    δs = [0.2, 0.21]
-    gs = [0.0 0.01; 0.01 0.0]
-
-    comp = MultiTransmonSystem(ωs, δs, gs)
-    @test comp isa CompositeQuantumSystem
-    @test length(comp.subsystems) == 2
-    @test !iszero(comp.H(zeros(comp.n_drives), 0.0))
-
-    comp2 = MultiTransmonSystem(
-        ωs, δs, gs;
-        levels_per_transmon=4,
-        subsystem_levels=[4,4],
-        subsystems=[1],
-        subsystem_drive_indices=[1]
-    )
-    @test comp2 isa CompositeQuantumSystem
-    @test length(comp2.subsystems) == 1
-    @test !isapprox(norm(comp2.H(zeros(comp2.n_drives), 0.0)), 0.0; atol=1e-12)
-end
-
-@testitem "MultiTransmonSystem: edge cases" begin
-    using PiccoloQuantumObjects
-    ωs = [4.0, 4.1, 4.2]
-    δs = [0.2, 0.21, 0.22]
-    gs = [0.0 0.01 0.02; 0.01 0.0 0.03; 0.02 0.03 0.0]
-    # Only a subset of subsystems
-    comp = MultiTransmonSystem(
-        ωs, δs, gs;
-        subsystems=[1,3],
-        subsystem_drive_indices=[3]
-    )
-    @test comp isa CompositeQuantumSystem
-    @test length(comp.subsystems) == 2
-    # Only one drive
-    @test comp.subsystems[1].n_drives == 0
-    @test comp.subsystems[2].n_drives == 2
-end
diff --git a/src/trajectory_initialization.jl b/src/trajectory_initialization.jl
deleted file mode 100644
index d5297064..00000000
--- a/src/trajectory_initialization.jl
+++ /dev/null
@@ -1,745 +0,0 @@
-module TrajectoryInitialization
-
-export unitary_geodesic
-export linear_interpolation
-export unitary_linear_interpolation
-export initialize_trajectory
-
-using NamedTrajectories
-import NamedTrajectories.StructNamedTrajectory: ScalarBound, VectorBound
-using PiccoloQuantumObjects
-
-using Distributions
-using ExponentialAction
-using LinearAlgebra
-using TestItems
-
-
-# ----------------------------------------------------------------------------- #
-#                           Initial states                                      #
-# ----------------------------------------------------------------------------- #
-
-linear_interpolation(x::AbstractVector, y::AbstractVector, n::Int) = hcat(range(x, y, n)...)
-
-"""
-    unitary_linear_interpolation(
-        U_init::AbstractMatrix,
-        U_goal::AbstractMatrix,
-        samples::Int
-    )
-
-Compute a linear interpolation of unitary operators with `samples` samples.
-"""
-function unitary_linear_interpolation(
-    U_init::AbstractMatrix{<:Number},
-    U_goal::AbstractMatrix{<:Number},
-    samples::Int
-)
-    Ũ⃗_init = operator_to_iso_vec(U_init)
-    Ũ⃗_goal = operator_to_iso_vec(U_goal)
-    Ũ⃗s = [Ũ⃗_init + (Ũ⃗_goal - Ũ⃗_init) * t for t ∈ range(0, 1, length=samples)]
-    Ũ⃗ = hcat(Ũ⃗s...)
-    return Ũ⃗
-end
-
-function unitary_linear_interpolation(
-    U_init::AbstractMatrix{<:Number},
-    U_goal::EmbeddedOperator,
-    samples::Int
-)
-    return unitary_linear_interpolation(U_init, U_goal.operator, samples)
-end
-
-"""
-    unitary_geodesic(U_init, U_goal, times; kwargs...)
-
-Compute the geodesic connecting U_init and U_goal at the specified times.
-
-# Arguments
-- `U_init::AbstractMatrix{<:Number}`: The initial unitary operator.
-- `U_goal::AbstractMatrix{<:Number}`: The goal unitary operator.
-- `times::AbstractVector{<:Number}`: The times at which to evaluate the geodesic.
-
-# Keyword Arguments
-- `return_unitary_isos::Bool=true`: If true returns a matrix where each column is a unitary
-    isovec, i.e. vec(vcat(real(U), imag(U))). If false, returns a vector of unitary matrices.
-- `return_generator::Bool=false`: If true, returns the effective Hamiltonian generating
-    the geodesic.
-"""
-function unitary_geodesic end
-
-function unitary_geodesic(
-    U_init::AbstractMatrix{<:Number},
-    U_goal::AbstractMatrix{<:Number},
-    times::AbstractVector{<:Number};
-    return_unitary_isos=true,
-    return_generator=false,
-    H_drift::AbstractMatrix{<:Number}=zeros(size(U_init)),
-)
-    t₀ = times[1]
-    T = times[end] - t₀
-
-    U_drift(t) = exp(-im * H_drift * t)
-    H = im * log(U_drift(T)' * (U_goal * U_init')) / T
-    # -im prefactor is not included in H
-    U_geo = [U_drift(t) * exp(-im * H * (t - t₀)) * U_init for t ∈ times]
-
-    if !return_unitary_isos
-        if return_generator
-            return U_geo, H
-        else
-            return U_geo
-        end
-    else
-        Ũ⃗_geo = stack(operator_to_iso_vec.(U_geo), dims=2)
-        if return_generator
-            return Ũ⃗_geo, H
-        else
-            return Ũ⃗_geo
-        end
-    end
-end
-
-function unitary_geodesic(
-    U_goal::AbstractPiccoloOperator,
-    samples::Int;
-    kwargs...
-)
-    return unitary_geodesic(
-        I(size(U_goal, 1)),
-        U_goal,
-        samples;
-        kwargs...
-    )
-end
-
-function unitary_geodesic(
-    U_init::AbstractMatrix{<:Number},
-    U_goal::EmbeddedOperator,
-    samples::Int;
-    H_drift::AbstractMatrix{<:Number}=zeros(size(U_init)),
-    kwargs...
-)
-    H_drift = unembed(H_drift, U_goal)
-    U1 = unembed(U_init, U_goal)
-    U2 = unembed(U_goal)
-    Ũ⃗ = unitary_geodesic(U1, U2, samples; H_drift=H_drift, kwargs...)
-    return hcat([
-        operator_to_iso_vec(embed(iso_vec_to_operator(Ũ⃗ₜ), U_goal))
-        for Ũ⃗ₜ ∈ eachcol(Ũ⃗)
-    ]...)
-end
-
-function unitary_geodesic(
-    U_init::AbstractMatrix{<:Number},
-    U_goal::AbstractMatrix{<:Number},
-    samples::Int;
-    kwargs...
-)
-    return unitary_geodesic(U_init, U_goal, range(0, 1, samples); kwargs...)
-end
-
-linear_interpolation(X::AbstractMatrix, Y::AbstractMatrix, n::Int) =
-    hcat([X + (Y - X) * t for t in range(0, 1, length=n)]...)
-
-# ============================================================================= #
-
-function initialize_unitary_trajectory(
-    U_init::AbstractMatrix{<:Number},
-    U_goal::AbstractPiccoloOperator,
-    N::Int;
-    geodesic::Bool=true,
-    system::Union{AbstractQuantumSystem, Nothing}=nothing
-)
-    if geodesic
-        if system isa AbstractQuantumSystem
-            H_drift = Matrix(get_drift(system))
-        else
-            H_drift = zeros(size(U_init))
-        end
-        Ũ⃗ = unitary_geodesic(U_init, U_goal, N, H_drift=H_drift)
-    else
-        Ũ⃗ = unitary_linear_interpolation(U_init, U_goal, N)
-    end
-    return Ũ⃗
-end
-
-# ----------------------------------------------------------------------------- #
-#                           Initial controls                                    #
-# ----------------------------------------------------------------------------- #
-
-function initialize_control_trajectory(
-    n_drives::Int,
-    n_derivatives::Int,
-    N::Int,
-    bounds::VectorBound,
-    drive_derivative_σ::Float64,
-)
-    if bounds isa AbstractVector
-        u_dists = [Uniform(-bounds[i], bounds[i]) for i = 1:n_drives]
-    elseif bounds isa Tuple
-        u_dists = [Uniform(uᵢ_lb, uᵢ_ub) for (uᵢ_lb, uᵢ_ub) ∈ zip(bounds...)]
-    else
-        error("bounds must be a Vector or Tuple")
-    end
-
-    controls = Matrix{Float64}[]
-
-    u = hcat([
-        zeros(n_drives),
-        vcat([rand(u_dists[i], 1, N - 2) for i = 1:n_drives]...),
-        zeros(n_drives)
-    ]...)
-    push!(controls, u)
-
-    for _ in 1:n_derivatives
-        push!(controls, randn(n_drives, N) * drive_derivative_σ)
-    end
-
-    return controls
-end
-
-function initialize_control_trajectory(
-    u::AbstractMatrix,
-    Δt::AbstractVecOrMat,
-    n_derivatives::Int
-)
-    controls = Matrix{Float64}[u]
-
-    for n in 1:n_derivatives
-        # next derivative
-        push!(controls, derivative(controls[end], Δt))
-
-        # to avoid constraint violation error at initial iteration for du, ddu, ...
-        if n > 1
-            controls[end-1][:, end] =
-                controls[end-1][:, end-1] + Δt[end-1] * controls[end][:, end-1]
-        end
-    end
-    return controls
-end
-
-initialize_control_trajectory(u::AbstractMatrix, Δt::Real, n_derivatives::Int) =
-    initialize_control_trajectory(u, fill(Δt, size(u, 2)), n_derivatives)
-
-# ----------------------------------------------------------------------------- #
-#                           Trajectory initialization                           #
-# ----------------------------------------------------------------------------- #
-
-"""
-    initialize_trajectory
-
-
-Initialize a trajectory for a control problem. The trajectory is initialized with
-data that should be consistently the same type (in this case, Float64).
-
-"""
-function initialize_trajectory(
-    state_data::Vector{<:AbstractMatrix{Float64}},
-    state_inits::Vector{<:AbstractVector{Float64}},
-    state_goals::Vector{<:AbstractVector{Float64}},
-    state_names::AbstractVector{Symbol},
-    N::Int,
-    Δt::Union{Float64, AbstractVecOrMat{<:Float64}},
-    n_drives::Int,
-    control_bounds::Tuple;
-    bound_state=false,
-    control_name=:u,
-    n_control_derivatives::Int=length(control_bounds) - 1,
-    zero_initial_and_final_derivative=false,
-    timestep_name=:Δt,
-    Δt_bounds::ScalarBound=(0.5 * Δt, 1.5 * Δt),
-    drive_derivative_σ::Float64=0.1,
-    u_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
-    global_component_data::NamedTuple{gname, <:Tuple{Vararg{AbstractVector{<:Real}}}} where gname=(;),
-    verbose=false,
-)
-    @assert length(state_data) == length(state_names) == length(state_inits) == length(state_goals) "state_data, state_names, state_inits, and state_goals must have the same length"
-    @assert length(control_bounds) == n_control_derivatives + 1 "control_bounds must have $n_control_derivatives + 1 elements"
-
-    # assert that state names are unique
-    @assert length(state_names) == length(Set(state_names)) "state_names must be unique"
-
-    # Control data
-    control_derivative_names = [
-        Symbol("d"^i * string(control_name)) for i = 1:n_control_derivatives
-    ]
-    if verbose
-        println("\tcontrol derivative names: $control_derivative_names")
-    end
-
-    control_names = (control_name, control_derivative_names...)
-
-    # Convert control_bounds elements from Vector{Tuple{Float64, Float64}} format to Tuple{Vector, Vector} format
-    # if necessary (to match NamedTrajectories expected types)
-    control_bounds_converted = map(control_bounds) do bound
-        if bound isa Vector{<:Tuple{Real, Real}}
-            # Convert from [(-1.5, 1.0), (-1.0, 2.0), ...] to ([lower...], [upper...])
-            return ([b[1] for b in bound], [b[2] for b in bound])
-        else
-            return bound
-        end
-    end
-
-    control_bounds = NamedTuple{control_names}(control_bounds_converted)
-
-    # Timestep data
-    if Δt isa Real
-        timestep_data = fill(Δt, 1, N)
-    elseif Δt isa AbstractVector
-        timestep_data = reshape(Δt, 1, :)
-    else
-        timestep_data = Δt
-        @assert size(Δt) == (1, N) "Δt must be a Real, AbstractVector, or 1x$(N) AbstractMatrix"
-    end
-    timestep = timestep_name
-
-    # Constraints
-    initial = (;
-        (state_names .=> state_inits)...,
-        control_name => zeros(n_drives),
-    )
-
-    final = (;
-        control_name => zeros(n_drives),
-    )
-
-    if zero_initial_and_final_derivative
-        initial = merge(initial, (; control_derivative_names[1] => zeros(n_drives),))
-        final = merge(final, (; control_derivative_names[1] => zeros(n_drives),))
-    end
-
-    goal = (; (state_names .=> state_goals)...)
-
-    # Bounds
-    bounds = control_bounds
-
-    bounds = merge(bounds, (; timestep_name => Δt_bounds,))
-
-    # Put unit box bounds on the state if bound_state is true
-    if bound_state
-        state_dim = length(state_inits[1])
-        state_bounds = repeat([(-ones(state_dim), ones(state_dim))], length(state_names))
-        bounds = merge(bounds, (; (state_names .=> state_bounds)...))
-    end
-
-    # Trajectory
-    if isnothing(u_guess)
-        # Randomly sample controls
-        control_data = initialize_control_trajectory(
-            n_drives,
-            n_control_derivatives,
-            N,
-            bounds[control_name],
-            drive_derivative_σ
-        )
-    else
-        # Use provided controls and take derivatives
-        control_data = initialize_control_trajectory(u_guess, Δt, n_control_derivatives)
-    end
-
-    names = [state_names..., control_names..., timestep_name]
-    values = [state_data..., control_data..., timestep_data]
-    controls = (control_names[end], timestep_name)
-
-    return NamedTrajectory(
-        (; (names .=> values)...),
-        global_component_data;
-        controls=controls,
-        timestep=timestep,
-        bounds=bounds,
-        initial=initial,
-        final=final,
-        goal=goal,
-    )
-end
-
-"""
-    initialize_trajectory
-
-Trajectory initialization of unitaries.
-"""
-function initialize_trajectory(
-    U_goal::AbstractPiccoloOperator,
-    N::Int,
-    Δt::Union{Real, AbstractVecOrMat{<:Real}},
-    args...;
-    state_name::Symbol=:Ũ⃗,
-    U_init::AbstractMatrix{<:Number}=Matrix{ComplexF64}(I(size(U_goal, 1))),
-    u_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
-    system::Union{AbstractQuantumSystem, Nothing}=nothing,
-    rollout_integrator::Function=expv,
-    geodesic=true,
-    kwargs...
-)
-    # Construct timesteps
-    if Δt isa AbstractMatrix
-        timesteps = vec(Δt)
-    elseif Δt isa Float64
-        timesteps = fill(Δt, N)
-    else
-        timesteps = Δt
-    end
-
-    # Initial state and goal
-    Ũ⃗_init = operator_to_iso_vec(U_init)
-
-    if U_goal isa EmbeddedOperator
-        Ũ⃗_goal = operator_to_iso_vec(U_goal.operator)
-    else
-        Ũ⃗_goal = operator_to_iso_vec(U_goal)
-    end
-
-    # Construct state data
-    if isnothing(u_guess)
-        Ũ⃗_traj = initialize_unitary_trajectory(
-            U_init,
-            U_goal,
-            N;
-            geodesic=geodesic,
-            system=system
-        )
-    else
-        @assert !isnothing(system) "System must be provided if u_guess is provided."
-        Ũ⃗_traj = unitary_rollout(Ũ⃗_init, u_guess, timesteps, system; integrator=rollout_integrator)
-    end
-
-    return initialize_trajectory(
-        [Ũ⃗_traj],
-        [Ũ⃗_init],
-        [Ũ⃗_goal],
-        [state_name],
-        N,
-        Δt,
-        args...;
-        u_guess=u_guess,
-        kwargs...
-    )
-end
-
-
-
-"""
-    initialize_trajectory
-
-Trajectory initialization of quantum states.
-"""
-function initialize_trajectory(
-    ψ_goals::AbstractVector{<:AbstractVector{ComplexF64}},
-    ψ_inits::AbstractVector{<:AbstractVector{ComplexF64}},
-    N::Int,
-    Δt::Union{Real, AbstractVector{<:Real}},
-    args...;
-    state_name=:ψ̃,
-    state_names::AbstractVector{<:Symbol}=length(ψ_goals) == 1 ?
-        [state_name] :
-        [Symbol(string(state_name) * "$i") for i = 1:length(ψ_goals)],
-    u_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
-    system::Union{AbstractQuantumSystem, Nothing}=nothing,
-    rollout_integrator::Function=expv,
-    kwargs...
-)
-    @assert length(ψ_inits) == length(ψ_goals) "ψ_inits and ψ_goals must have the same length"
-    @assert length(state_names) == length(ψ_goals) "state_names and ψ_goals must have the same length"
-
-    ψ̃_goals = ket_to_iso.(ψ_goals)
-    ψ̃_inits = ket_to_iso.(ψ_inits)
-
-    # Construct timesteps
-    if Δt isa AbstractMatrix
-        timesteps = vec(Δt)
-    elseif Δt isa Float64
-        timesteps = fill(Δt, N)
-    else
-        timesteps = Δt
-    end
-
-    # Construct state data
-    ψ̃_trajs = Matrix{Float64}[]
-    if isnothing(u_guess)
-        for (ψ̃_init, ψ̃_goal) ∈ zip(ψ̃_inits, ψ̃_goals)
-            ψ̃_traj = linear_interpolation(ψ̃_init, ψ̃_goal, N)
-            push!(ψ̃_trajs, ψ̃_traj)
-        end
-        if system isa AbstractVector
-            ψ̃_trajs = repeat(ψ̃_trajs, length(system))
-        end
-    else
-        for ψ̃_init ∈ ψ̃_inits
-            ψ̃_traj = rollout(ψ̃_init, u_guess, timesteps, system; integrator=rollout_integrator)
-            push!(ψ̃_trajs, ψ̃_traj)
-        end
-    end
-
-    return initialize_trajectory(
-        ψ̃_trajs,
-        ψ̃_inits,
-        ψ̃_goals,
-        state_names,
-        N,
-        Δt,
-        args...;
-        u_guess=u_guess,
-        kwargs...
-    )
-end
-
-"""
-    initialize_trajectory
-
-Trajectory initialization of density matrices.
-"""
-function initialize_trajectory(
-    ρ_init,
-    ρ_goal,
-    N::Int,
-    Δt::Union{Real, AbstractVecOrMat{<:Real}},
-    args...;
-    state_name::Symbol=:ρ⃗̃,
-    u_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
-    system::Union{OpenQuantumSystem, Nothing}=nothing,
-    rollout_integrator::Function=expv,
-    kwargs...
-)
-    # Construct timesteps
-    if Δt isa AbstractMatrix
-        timesteps = vec(Δt)
-    elseif Δt isa Float64
-        timesteps = fill(Δt, N)
-    else
-        timesteps = Δt
-    end
-
-    # Initial state and goal
-    ρ⃗̃_init = density_to_iso_vec(ρ_init)
-    ρ⃗̃_goal = density_to_iso_vec(ρ_goal)
-
-    # Construct state data
-    if isnothing(u_guess)
-        ρ⃗̃_traj = linear_interpolation(ρ_init, ρ_goal, N)
-    else
-        @assert !isnothing(system) "System must be provided if u_guess is provided."
-
-        ρ⃗̃_traj = open_rollout(
-            ρ_init,
-            u_guess,
-            timesteps,
-            system;
-            integrator=rollout_integrator
-        )
-    end
-
-    return initialize_trajectory(
-        [ρ⃗̃_traj],
-        [ρ⃗̃_init],
-        [ρ⃗̃_goal],
-        [state_name],
-        N,
-        Δt,
-        args...;
-        u_guess=u_guess,
-        kwargs...
-    )
-end
-
-
-
-# ============================================================================= #
-
-@testitem "Random drive initialization" begin
-    T = 10
-    n_drives = 2
-    n_derivates = 2
-    drive_bounds = [1.0, 2.0]
-    drive_derivative_σ = 0.01
-
-    u, du, ddu = TrajectoryInitialization.initialize_control_trajectory(n_drives, n_derivates, T, drive_bounds, drive_derivative_σ)
-
-    @test size(u) == (n_drives, T)
-    @test size(du) == (n_drives, T)
-    @test size(ddu) == (n_drives, T)
-    @test all([-drive_bounds[i] < minimum(u[i, :]) < drive_bounds[i] for i in 1:n_drives])
-end
-
-@testitem "Geodesic" begin
-    using LinearAlgebra
-    using PiccoloQuantumObjects
-
-    ## Group 1: identity to X (π rotation)
-
-    # Test π rotation
-    U_α = GATES[:I]
-    U_ω = GATES[:X]
-    Us, H = unitary_geodesic(
-        U_α, U_ω, range(0, 1, 4), return_generator=true
-    )
-
-    @test size(Us, 2) == 4
-    @test Us[:, 1] ≈ operator_to_iso_vec(U_α)
-    @test Us[:, end] ≈ operator_to_iso_vec(U_ω)
-    @test H' - H ≈ zeros(2, 2)
-    @test norm(H) ≈ π
-
-    # Test modified timesteps (10x)
-    Us10, H10 = unitary_geodesic(
-        U_α, U_ω, range(-5, 5, 4), return_generator=true
-    )
-
-    @test size(Us10, 2) == 4
-    @test Us10[:, 1] ≈ operator_to_iso_vec(U_α)
-    @test Us10[:, end] ≈ operator_to_iso_vec(U_ω)
-    @test H10' - H10 ≈ zeros(2, 2)
-    @test norm(H10) ≈ π/10
-
-    # Test wrapped call
-    Us_wrap, H_wrap = unitary_geodesic(U_ω, 10, return_generator=true)
-    @test Us_wrap[:, 1] ≈ operator_to_iso_vec(GATES[:I])
-    @test Us_wrap[:, end] ≈ operator_to_iso_vec(U_ω)
-    rotation = [exp(-im * H_wrap * t) for t ∈ range(0, 1, 10)]
-    Us_test = stack(operator_to_iso_vec.(rotation), dims=2)
-    @test isapprox(Us_wrap, Us_test)
-
-
-    ## Group 2: √X to X (π/2 rotation)
-
-    # Test geodesic not at identity
-    U₀ = sqrt(GATES[:X])
-    U₁ = GATES[:X]
-    Us, H = unitary_geodesic(U₀, U₁, 10, return_generator=true)
-    @test Us[:, 1] ≈ operator_to_iso_vec(U₀)
-    @test Us[:, end] ≈ operator_to_iso_vec(U_ω)
-
-    rotation = [exp(-im * H * t) * U₀ for t ∈ range(0, 1, 10)]
-    Us_test = stack(operator_to_iso_vec.(rotation), dims=2)
-    @test isapprox(Us, Us_test)
-    Us_wrap = unitary_geodesic(U_ω, 4)
-    @test Us_wrap[:, 1] ≈ operator_to_iso_vec(GATES[:I])
-    @test Us_wrap[:, end] ≈ operator_to_iso_vec(U_ω)
-
-end
-
-@testitem "unitary trajectory initialization" begin
-    using NamedTrajectories
-    using PiccoloQuantumObjects 
-
-    U_goal = GATES[:X]
-    T = 10
-    Δt = 0.1
-    n_drives = 2
-    u_bounds = ([1.0, 1.0],)
-
-    traj = initialize_trajectory(
-        U_goal, T, Δt, n_drives, u_bounds
-    )
-
-    @test traj isa NamedTrajectory
-end
-@testitem "quantum state trajectory initialization" begin
-    using NamedTrajectories
-
-    ψ_init = Vector{ComplexF64}([0.0, 1.0])
-    ψ_goal = Vector{ComplexF64}([1.0, 0.0])
-
-    T = 10
-    Δt = 0.1
-    n_drives = 2
-    all_u_bounds = ([1.0, 1.0],)
-
-    traj = initialize_trajectory(
-        [ψ_goal], [ψ_init], T, Δt, n_drives, all_u_bounds
-    )
-
-    @test traj isa NamedTrajectory
-end
-
-@testitem "unitary_linear_interpolation direct" begin
-    using PiccoloQuantumObjects
-    U_init = GATES[:I]
-    U_goal = GATES[:X]
-    samples = 5
-    # Direct matrix
-    Ũ⃗ = TrajectoryInitialization.unitary_linear_interpolation(U_init, U_goal, samples)
-    @test size(Ũ⃗, 2) == samples
-    # EmbeddedOperator
-    U_init_emb = EmbeddedOperator(U_init, [1,2], [2,2])
-    U_goal_emb = EmbeddedOperator(U_goal, [1,2], [2,2])
-    Ũ⃗2 = TrajectoryInitialization.unitary_linear_interpolation(U_init_emb.operator, U_goal_emb, samples)
-    @test size(Ũ⃗2, 2) == samples
-end
-
-@testitem "initialize_unitary_trajectory geodesic=false" begin
-    using PiccoloQuantumObjects
-    U_init = GATES[:I]
-    U_goal = GATES[:X]
-    T = 4
-    Ũ⃗ = TrajectoryInitialization.initialize_unitary_trajectory(U_init, U_goal, T; geodesic=false)
-    @test size(Ũ⃗, 2) == T
-end
-
-@testitem "initialize_control_trajectory with a, Δt, n_derivatives" begin
-    n_drives = 2
-    T = 5
-    n_derivatives = 2
-    u = randn(n_drives, T)
-    Δt = fill(0.1, T)
-    controls = TrajectoryInitialization.initialize_control_trajectory(u, Δt, n_derivatives)
-    @test length(controls) == n_derivatives + 1
-    @test size(controls[1]) == (n_drives, T)
-    # Real Δt version
-    controls2 = TrajectoryInitialization.initialize_control_trajectory(u, 0.1, n_derivatives)
-    @test length(controls2) == n_derivatives + 1
-end
-
-@testitem "initialize_trajectory with bound_state and zero_initial_and_final_derivative" begin
-    using NamedTrajectories: NamedTrajectory
-    state_data = [rand(2, 4)]
-    state_inits = [rand(2)]
-    state_goals = [rand(2)]
-    state_names = [:x]
-    T = 4
-    Δt = 0.1
-    n_drives = 1
-    control_bounds = ([1.0], [1.0])
-    traj = TrajectoryInitialization.initialize_trajectory(
-        state_data, state_inits, state_goals, state_names, T, Δt, n_drives, control_bounds;
-        bound_state=true, zero_initial_and_final_derivative=true
-    )
-    @test traj isa NamedTrajectory
-end
-
-@testitem "initialize_trajectory error branches" begin
-    state_data = [rand(2, 4)]
-    state_inits = [rand(2)]
-    state_goals = [rand(2)]
-    state_names = [:x]
-    T = 4
-    Δt = 0.1
-    n_drives = 1
-    control_bounds = ([1.0], [1.0])
-    # state_names not unique
-    @test_throws AssertionError TrajectoryInitialization.initialize_trajectory(
-        state_data, state_inits, state_goals, [:x, :x], T, Δt, n_drives, control_bounds
-    )
-    # control_bounds wrong length
-    @test_throws AssertionError TrajectoryInitialization.initialize_trajectory(
-        state_data, state_inits, state_goals, state_names, T, Δt, n_drives, ([1.0],); n_control_derivatives=1
-    )
-    # bounds wrong type
-    @test_throws MethodError TrajectoryInitialization.initialize_control_trajectory(
-        n_drives, 2, T, "notabounds", 0.1
-    )
-end
-
-@testitem "linear_interpolation for matrices" begin
-    X = [1.0 2.0; 3.0 4.0]
-    Y = [5.0 6.0; 7.0 8.0]
-    n = 3
-    result = linear_interpolation(X, Y, n)
-    @test size(result) == (2, 2 * n)
-    @test result[:, 1:2] ≈ X
-    @test result[:, 5:6] ≈ Y
-    @test result[:, 3:4] ≈ (X + Y) / 2
-end
-
-
-end
diff --git a/test/runtests.jl b/test/runtests.jl
index 61d982c9..e7060bb1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,4 +1,4 @@
 using TestItemRunner
 
 # Run all testitem tests in package
-@run_package_tests
+@run_package_tests filter=ti->( !(:experimental in ti.tags) )
diff --git a/test/scripts/test_changes.jl b/test/scripts/test_changes.jl
new file mode 100644
index 00000000..5c843b62
--- /dev/null
+++ b/test/scripts/test_changes.jl
@@ -0,0 +1,59 @@
+using QuantumCollocation
+using PiccoloQuantumObjects
+
+println("Testing basic system creation...")
+try
+    sys = TransmonSystem(10.0, [0.2, 0.2]; ω=4.0, δ=0.2, levels=3)
+    println("✓ TransmonSystem created successfully")
+    println("  n_drives: ", sys.n_drives)
+    println("  levels: ", sys.levels)
+catch e
+    println("✗ Error creating TransmonSystem:")
+    println(e)
+end
+
+println("\nTesting QuantumStateSmoothPulseProblem...")
+try
+    N = 51
+    Δt = 0.2
+    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+    ψ_init = Vector{ComplexF64}([1.0, 0.0])
+    ψ_target = Vector{ComplexF64}([0.0, 1.0])
+    
+    prob = QuantumStateSmoothPulseProblem(
+        sys, ψ_init, ψ_target, N, Δt;
+        piccolo_options=PiccoloOptions(verbose=false)
+    )
+    println("✓ QuantumStateSmoothPulseProblem created successfully")
+    println("  trajectory N: ", prob.trajectory.N)
+catch e
+    println("✗ Error creating problem:")
+    println(e)
+    for (exc, bt) in Base.catch_stack()
+        showerror(stdout, exc, bt)
+        println()
+    end
+end
+
+println("\nTesting UnitarySmoothPulseProblem...")
+try
+    N = 51
+    sys = TransmonSystem(10.0, [0.2, 0.2]; ω=4.0, δ=0.2, levels=2)
+    U_goal = GATES[:X]
+    
+    prob = UnitarySmoothPulseProblem(
+        sys, U_goal, N;
+        piccolo_options=PiccoloOptions(verbose=false)
+    )
+    println("✓ UnitarySmoothPulseProblem created successfully")
+    println("  trajectory N: ", prob.trajectory.N)
+catch e
+    println("✗ Error creating problem:")
+    println(e)
+    for (exc, bt) in Base.catch_stack()
+        showerror(stdout, exc, bt)
+        println()
+    end
+end
+
+println("\nAll basic tests completed!")
diff --git a/test/scripts/test_default_integrators.jl b/test/scripts/test_default_integrators.jl
new file mode 100644
index 00000000..f97ac199
--- /dev/null
+++ b/test/scripts/test_default_integrators.jl
@@ -0,0 +1,109 @@
+using QuantumCollocation
+using PiccoloQuantumObjects
+
+println("=" ^ 70)
+println("Testing Default Integrators")
+println("=" ^ 70)
+println()
+
+# Create a simple quantum system
+sys = QuantumSystem(
+    GATES[:Z],              # H_drift
+    [GATES[:X], GATES[:Y]], # H_drives
+    1.0,                    # T_max
+    [1.0, 1.0]             # drive_bounds
+)
+
+# ----------------------------------------------------------------------------- #
+# Example 1: Using default integrator
+# ----------------------------------------------------------------------------- #
+
+println("Example 1: Default integrator (simplest approach)")
+println("-" ^ 70)
+
+qtraj = unitary_trajectory(sys, GATES[:H], 10)
+println("✓ Created UnitaryTrajectory")
+
+# Get default integrator - no need to specify anything!
+integrator = default_integrator(qtraj)
+println("✓ Created default integrator")
+println("  Type: ", typeof(integrator))
+println()
+
+# ----------------------------------------------------------------------------- #
+# Example 2: Creating custom integrators when needed
+# ----------------------------------------------------------------------------- #
+
+println("Example 2: Custom integrator (when user wants control)")
+println("-" ^ 70)
+
+qtraj2 = unitary_trajectory(sys, GATES[:X], 10)
+println("✓ Created UnitaryTrajectory")
+
+# User can create custom integrator using trajectory info
+custom_integrator = UnitaryIntegrator(
+    get_system(qtraj2),
+    get_trajectory(qtraj2),
+    get_state_name(qtraj2),
+    get_control_name(qtraj2)
+)
+println("✓ Created custom UnitaryIntegrator")
+println("  Type: ", typeof(custom_integrator))
+println()
+
+# ----------------------------------------------------------------------------- #
+# Example 3: Adding derivative constraints
+# ----------------------------------------------------------------------------- #
+
+println("Example 3: Combining with derivative integrators")
+println("-" ^ 70)
+
+qtraj3 = unitary_trajectory(sys, GATES[:Y], 10)
+println("✓ Created UnitaryTrajectory")
+
+# Get default dynamics integrator
+dynamics = default_integrator(qtraj3)
+println("✓ Created dynamics integrator")
+
+# Add derivative constraints directly using underlying trajectory
+using DirectTrajOpt
+du_integrator = DerivativeIntegrator(get_trajectory(qtraj3), :u, :du)
+ddu_integrator = DerivativeIntegrator(get_trajectory(qtraj3), :du, :ddu)
+println("✓ Created derivative integrators")
+
+# Combine all integrators
+all_integrators = [dynamics, du_integrator, ddu_integrator]
+println("✓ Combined ", length(all_integrators), " integrators")
+println()
+
+# ----------------------------------------------------------------------------- #
+# Example 4: Works for all trajectory types
+# ----------------------------------------------------------------------------- #
+
+println("Example 4: Default integrators for all trajectory types")
+println("-" ^ 70)
+
+# Ket trajectory
+ψ_init = ComplexF64[1.0, 0.0]
+ψ_goal = ComplexF64[0.0, 1.0]
+ket_traj = ket_trajectory(sys, ψ_init, ψ_goal, 10)
+ket_int = default_integrator(ket_traj)
+println("✓ KetTrajectory → ", typeof(ket_int))
+
+# Density trajectory
+open_sys = OpenQuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]], 1.0, [1.0, 1.0])
+ρ_init = ComplexF64[1.0 0.0; 0.0 0.0]
+ρ_goal = ComplexF64[0.0 0.0; 0.0 1.0]
+density_traj = density_trajectory(open_sys, ρ_init, ρ_goal, 10)
+density_int = default_integrator(density_traj)
+println("✓ DensityTrajectory → ", typeof(density_int))
+println()
+
+println("=" ^ 70)
+println("All examples completed successfully! ✓")
+println("=" ^ 70)
+println()
+println("Key Takeaway:")
+println("  • Use default_integrator(qtraj) for simple cases")
+println("  • Create custom integrators when you need control")
+println("  • Access trajectory info via get_system(), get_trajectory(), get_state_name(), etc.")
diff --git a/test/scripts/test_load.jl b/test/scripts/test_load.jl
new file mode 100644
index 00000000..572e6f01
--- /dev/null
+++ b/test/scripts/test_load.jl
@@ -0,0 +1,16 @@
+#!/usr/bin/env julia
+using Pkg
+Pkg.activate(".")
+try
+    @info "Attempting to load QuantumCollocation..."
+    using QuantumCollocation
+    println("✓ QuantumCollocation loaded successfully!")
+    println("✓ QuantumControlProblem available: ", isdefined(QuantumCollocation, :QuantumControlProblem))
+    println("✓ SmoothPulseProblem available: ", isdefined(QuantumCollocation, :SmoothPulseProblem))
+    println("✓ MinimumTimeProblem available: ", isdefined(QuantumCollocation, :MinimumTimeProblem))
+catch e
+    println("✗ Error loading QuantumCollocation:")
+    showerror(stdout, e, catch_backtrace())
+    println()
+    exit(1)
+end
diff --git a/test/scripts/test_quantum_trajectories.jl b/test/scripts/test_quantum_trajectories.jl
new file mode 100644
index 00000000..f0a6d74e
--- /dev/null
+++ b/test/scripts/test_quantum_trajectories.jl
@@ -0,0 +1,63 @@
+using QuantumCollocation
+using PiccoloQuantumObjects
+
+# Test UnitaryTrajectory
+sys = QuantumSystem(
+    GATES[:Z],              # H_drift
+    [GATES[:X], GATES[:Y]], # H_drives
+    1.0,                    # T_max
+    [1.0, 1.0]             # drive_bounds
+)
+
+U_goal = GATES[:H]  # Hadamard gate
+N = 10
+
+println("Testing unitary_trajectory...")
+qtraj = unitary_trajectory(sys, U_goal, N)
+println("  Type: ", typeof(qtraj))
+println("  Is UnitaryTrajectory: ", qtraj isa UnitaryTrajectory)
+println("  System: ", get_system(qtraj) === sys)
+println("  Goal: ", get_goal(qtraj) == U_goal)
+println("  State name: ", get_state_name(qtraj))
+println("  Control name: ", get_control_name(qtraj))
+println("  Can access bounds: ", haskey(qtraj.bounds, :u))
+println("  Can access components: ", haskey(qtraj.components, :Ũ⃗))
+println("  Can index: ", size(qtraj[:u]))
+println()
+
+# Test KetTrajectory
+ψ_init = ComplexF64[1.0, 0.0]
+ψ_goal = ComplexF64[0.0, 1.0]
+
+println("Testing ket_trajectory...")
+qtraj2 = ket_trajectory(sys, ψ_init, ψ_goal, N)
+println("  Type: ", typeof(qtraj2))
+println("  Is KetTrajectory: ", qtraj2 isa KetTrajectory)
+println("  System: ", get_system(qtraj2) === sys)
+println("  Goal: ", get_goal(qtraj2) == ψ_goal)
+println("  State name: ", get_state_name(qtraj2))
+println("  Control name: ", get_control_name(qtraj2))
+println()
+
+# Test DensityTrajectory
+open_sys = OpenQuantumSystem(
+    GATES[:Z],              # H_drift
+    [GATES[:X], GATES[:Y]], # H_drives
+    1.0,                    # T_max
+    [1.0, 1.0]             # drive_bounds
+)
+
+ρ_init = ComplexF64[1.0 0.0; 0.0 0.0]
+ρ_goal = ComplexF64[0.0 0.0; 0.0 1.0]
+
+println("Testing density_trajectory...")
+qtraj3 = density_trajectory(open_sys, ρ_init, ρ_goal, N)
+println("  Type: ", typeof(qtraj3))
+println("  Is DensityTrajectory: ", qtraj3 isa DensityTrajectory)
+println("  System: ", get_system(qtraj3) === open_sys)
+println("  Goal: ", get_goal(qtraj3) == ρ_goal)
+println("  State name: ", get_state_name(qtraj3))
+println("  Control name: ", get_control_name(qtraj3))
+println()
+
+println("All basic tests passed! ✓")
diff --git a/test/scripts/test_smooth_pulse_problem.jl b/test/scripts/test_smooth_pulse_problem.jl
new file mode 100644
index 00000000..49e1a715
--- /dev/null
+++ b/test/scripts/test_smooth_pulse_problem.jl
@@ -0,0 +1,136 @@
+using QuantumCollocation
+using PiccoloQuantumObjects
+using Test
+
+println("=" ^ 70)
+println("Testing Unified SmoothPulseProblem")
+println("=" ^ 70)
+println()
+
+# =============================================================================
+# Test 1: UnitaryTrajectory → SmoothPulseProblem
+# =============================================================================
+
+println("Test 1: SmoothPulseProblem with UnitaryTrajectory")
+println("-" ^ 70)
+
+sys = QuantumSystem(
+    GATES[:Z],
+    [GATES[:X], GATES[:Y]],
+    1.0,
+    [1.0, 1.0]
+)
+
+# Step 1: Create quantum trajectory
+qtraj = unitary_trajectory(sys, GATES[:H], 10)
+println("✓ Created UnitaryTrajectory")
+
+# Step 2: Create smooth pulse problem (automatically adds derivatives!)
+prob = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)
+println("✓ Created SmoothPulseProblem")
+
+# Verify structure
+@test prob isa DirectTrajOptProblem
+println("✓ Problem is DirectTrajOptProblem")
+
+@test length(prob.integrators) == 3  # dynamics + du + ddu
+println("✓ Has 3 integrators (dynamics, du, ddu)")
+
+@test haskey(prob.trajectory.components, :u)
+@test haskey(prob.trajectory.components, :du)
+@test haskey(prob.trajectory.components, :ddu)
+println("✓ Trajectory has u, du, ddu components")
+println()
+
+# =============================================================================
+# Test 2: KetTrajectory → SmoothPulseProblem
+# =============================================================================
+
+println("Test 2: SmoothPulseProblem with KetTrajectory")
+println("-" ^ 70)
+
+ψ_init = ComplexF64[1.0, 0.0]
+ψ_goal = ComplexF64[0.0, 1.0]
+
+qtraj_ket = ket_trajectory(sys, ψ_init, ψ_goal, 10)
+println("✓ Created KetTrajectory")
+
+prob_ket = SmoothPulseProblem(qtraj_ket; Q=50.0, R=1e-3)
+println("✓ Created SmoothPulseProblem")
+
+@test prob_ket isa DirectTrajOptProblem
+@test length(prob_ket.integrators) == 3
+println("✓ Structure correct")
+println()
+
+# =============================================================================
+# Test 3: DensityTrajectory → SmoothPulseProblem
+# =============================================================================
+
+println("Test 3: SmoothPulseProblem with DensityTrajectory")
+println("-" ^ 70)
+
+open_sys = OpenQuantumSystem(
+    GATES[:Z],
+    [GATES[:X], GATES[:Y]],
+    1.0,
+    [1.0, 1.0]
+)
+
+ρ_init = ComplexF64[1.0 0.0; 0.0 0.0]
+ρ_goal = ComplexF64[0.0 0.0; 0.0 1.0]
+
+qtraj_density = density_trajectory(open_sys, ρ_init, ρ_goal, 10)
+println("✓ Created DensityTrajectory")
+
+prob_density = SmoothPulseProblem(qtraj_density; Q=100.0)
+println("✓ Created SmoothPulseProblem")
+
+@test prob_density isa DirectTrajOptProblem
+@test length(prob_density.integrators) == 3
+println("✓ Structure correct")
+println()
+
+# =============================================================================
+# Test 4: Custom n_derivatives
+# =============================================================================
+
+println("Test 4: Custom number of derivatives")
+println("-" ^ 70)
+
+qtraj1 = unitary_trajectory(sys, GATES[:X], 10)
+
+# Only first derivative
+prob_1deriv = SmoothPulseProblem(qtraj1; n_derivatives=1)
+@test length(prob_1deriv.integrators) == 2  # dynamics + du
+@test haskey(prob_1deriv.trajectory.components, :du)
+@test !haskey(prob_1deriv.trajectory.components, :ddu)
+println("✓ n_derivatives=1 works (2 integrators)")
+
+# No derivatives
+qtraj2 = unitary_trajectory(sys, GATES[:Y], 10)
+prob_0deriv = SmoothPulseProblem(qtraj2; n_derivatives=0)
+@test length(prob_0deriv.integrators) == 1  # only dynamics
+@test !haskey(prob_0deriv.trajectory.components, :du)
+println("✓ n_derivatives=0 works (1 integrator)")
+println()
+
+# =============================================================================
+# Summary
+# =============================================================================
+
+println("=" ^ 70)
+println("All tests passed! ✓")
+println("=" ^ 70)
+println()
+println("Key Features Demonstrated:")
+println("  1. Unified interface: SmoothPulseProblem(qtraj)")
+println("  2. Type dispatch: works for Unitary/Ket/Density trajectories")
+println("  3. Automatic derivative addition via add_control_derivatives()")
+println("  4. Default integrators via default_integrator()")
+println("  5. Flexible n_derivatives parameter")
+println()
+println("Usage Pattern:")
+println("  qtraj = unitary_trajectory(sys, U_goal, N)")
+println("  prob = SmoothPulseProblem(qtraj; Q=100.0, R=1e-2)")
+println("  solve!(prob; max_iter=100)")