diff --git a/src/problem_templates/quantum_state_sampling_problem.jl b/src/problem_templates/quantum_state_sampling_problem.jl
index 6d8e47d2..705a5a55 100644
--- a/src/problem_templates/quantum_state_sampling_problem.jl
+++ b/src/problem_templates/quantum_state_sampling_problem.jl
@@ -160,13 +160,25 @@ end
         piccolo_options=PiccoloOptions(verbose=false)
     )
 
-    state_names = [n for n ∈ prob.trajectory.names if startswith(string(n), "ψ̃")]
+    state_name = :ψ̃
+    state_names = [n for n ∈ prob.trajectory.names if startswith(string(n), string(state_name))]
+    sys1_state_names = [n for n ∈ state_names if endswith(string(n), "1")]
 
-    init = [fidelity(prob.trajectory, sys1, state_symb=n) for n in state_names]
+    # Separately compute all unique initial and goal state fidelities 
+    init = [fidelity(prob.trajectory, sys1, state_symb=n) for n in sys1_state_names]
     solve!(prob, max_iter=20)
-    final = [fidelity(prob.trajectory, sys1, state_symb=n) for n in state_names]
+    final = [fidelity(prob.trajectory, sys1, state_symb=n) for n in sys1_state_names]
     @test all(final .> init)
 
+    # Check that a_guess can be used
+    prob = QuantumStateSamplingProblem(
+        [sys1, sys2], ψ_init, ψ_target, T, Δt;
+        ipopt_options=IpoptOptions(print_level=1),
+        piccolo_options=PiccoloOptions(verbose=false),
+        a_guess=prob.trajectory.a
+    )
+    solve!(prob, max_iter=20)
+
     # Compare a solution without robustness
     # -------------------------------------
     prob_default = QuantumStateSmoothPulseProblem(
@@ -177,6 +189,7 @@ end
     )
     solve!(prob_default, max_iter=20)
     final_default = fidelity(prob_default.trajectory, sys1)
+    # Pick any initial state
     final_robust = fidelity(prob.trajectory, sys1, state_symb=state_names[1])
     @test final_robust > final_default
 end
@@ -199,11 +212,13 @@ end
         piccolo_options=PiccoloOptions(verbose=false)
     )
 
-    state_names = [n for n ∈ prob.trajectory.names if startswith(string(n), "ψ̃")]
+    state_name = :ψ̃
+    state_names = [n for n ∈ prob.trajectory.names if startswith(string(n), string(state_name))]
+    sys1_state_names = [n for n ∈ state_names if endswith(string(n), "1")]
 
-    init = [fidelity(prob.trajectory, sys1, state_symb=n) for n in state_names]
+    init = [fidelity(prob.trajectory, sys1, state_symb=n) for n in sys1_state_names]
     solve!(prob, max_iter=20)
-    final = [fidelity(prob.trajectory, sys1, state_symb=n) for n in state_names]
+    final = [fidelity(prob.trajectory, sys1, state_symb=n) for n in sys1_state_names]
     @test all(final .> init)
 end
 
diff --git a/src/trajectory_initialization.jl b/src/trajectory_initialization.jl
index f78f0bb3..5c97fa2a 100644
--- a/src/trajectory_initialization.jl
+++ b/src/trajectory_initialization.jl
@@ -513,11 +513,9 @@ function initialize_quantum_states(
             state_data = repeat(state_data, length(system))
         end
     else
-        for sys ∈ system
-            for ψ̃_init ∈ ψ̃_inits
-                ψ̃_traj = rollout(ψ̃_init, a_guess, timesteps, sys; integrator=rollout_integrator)
-                push!(state_data, ψ̃_traj)
-            end
+        for ψ̃_init ∈ ψ̃_inits
+            ψ̃_traj = rollout(ψ̃_init, a_guess, timesteps, system; integrator=rollout_integrator)
+            push!(state_data, ψ̃_traj)
         end
     end