diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 26f27990..5bf642de 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -34,9 +34,11 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v2
+      - uses: codecov/codecov-action@v3
         with:
           files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: false
   docs:
     name: Documentation
     runs-on: ubuntu-latest
diff --git a/Project.toml b/Project.toml
index 23d88e1d..4daae86f 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
 authors = ["Aaron Trowbridge <aaron.j.trowbridge@gmail.com> and contributors"]
-version = "0.6.0"
+version = "0.6.1"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -26,7 +26,7 @@ ExponentialAction = "0.2"
 Interpolations = "0.15"
 JLD2 = "0.5"
 LinearAlgebra = "1.10, 1.11"
-NamedTrajectories = "0.2"
+NamedTrajectories = "0.3"
 PiccoloQuantumObjects = "0.3"
 QuantumCollocationCore = "0.3"
 Random = "1.10, 1.11"
diff --git a/src/problem_templates/_problem_templates.jl b/src/problem_templates/_problem_templates.jl
index 9beda1e5..17a3745d 100644
--- a/src/problem_templates/_problem_templates.jl
+++ b/src/problem_templates/_problem_templates.jl
@@ -1,15 +1,13 @@
 module ProblemTemplates
 
-using ..DirectSums
-using ..Rollouts
 using ..TrajectoryInitialization
-using ..Losses
 
-using Distributions
 using TrajectoryIndexingUtils
 using NamedTrajectories
 using QuantumCollocationCore
 using PiccoloQuantumObjects
+
+using Distributions
 using LinearAlgebra
 using SparseArrays
 using ExponentialAction
diff --git a/src/problem_templates/unitary_direct_sum_problem.jl b/src/problem_templates/unitary_direct_sum_problem.jl
index 1bf5b7bb..aa7a086d 100644
--- a/src/problem_templates/unitary_direct_sum_problem.jl
+++ b/src/problem_templates/unitary_direct_sum_problem.jl
@@ -43,11 +43,11 @@ between each neighbor of the provided `probs`.
 
 """
 function UnitaryDirectSumProblem(
-    probs::AbstractVector{<:QuantumControlProblem},
+    probs::AbstractVector{QuantumControlProblem},
     final_fidelity::Real;
-    prob_labels::AbstractVector{<:String}=[string(i) for i ∈ 1:length(probs)],
+    prob_labels::AbstractVector{String}=[string(i) for i ∈ 1:length(probs)],
     graph::Union{Nothing, AbstractVector{<:Tuple{String, String}}, AbstractVector{<:Tuple{Symbol, Symbol}}}=nothing,
-    boundary_values::Union{AbstractDict{<:String, <:AbstractArray}, AbstractDict{<:Symbol, <:AbstractArray}}=Dict{String, Array}(),
+    boundary_values::Union{AbstractDict{String, <:AbstractArray}, AbstractDict{Symbol, <:AbstractArray}}=Dict{String, Array}(),
     control_name::Symbol=:a,
     Q::Union{Float64, Vector{Float64}}=100.0,
     Q_symb::Symbol=:dda,
@@ -59,7 +59,7 @@ function UnitaryDirectSumProblem(
     drive_derivative_σ::Float64=0.01,
     drive_reset_ratio::Float64=0.50,
     fidelity_cost::Bool=false,
-    subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
+    subspace::Union{AbstractVector{Int}, Nothing}=nothing,
     ipopt_options::IpoptOptions=deepcopy(probs[1].ipopt_options),
     piccolo_options::PiccoloOptions=deepcopy(probs[1].piccolo_options),
     kwargs...
@@ -185,7 +185,7 @@ end
 
 @testitem "Construct direct sum problem" begin
     using QuantumCollocationCore
-    using PiccoloQuantumObjects
+    using NamedTrajectories
     sys = QuantumSystem(0.01 * GATES[:Z], [GATES[:X], GATES[:Y]])
     U_goal1 = GATES[:X]
     U_ε = haar_identity(2, 0.33)
diff --git a/src/trajectory_initialization.jl b/src/trajectory_initialization.jl
index cb38de92..0d8c68fa 100644
--- a/src/trajectory_initialization.jl
+++ b/src/trajectory_initialization.jl
@@ -6,18 +6,13 @@ export unitary_linear_interpolation
 export initialize_trajectory
 
 using NamedTrajectories
+using PiccoloQuantumObjects
 
 using Distributions
 using ExponentialAction
 using LinearAlgebra
 using TestItems
 
-using ..Isomorphisms
-using ..QuantumSystems
-using ..Rollouts
-using ..EmbeddedOperators
-using ..DirectSums
-
 
 # ----------------------------------------------------------------------------- #
 #                           Initial states                                      #