Fix state dim methods

Author: bjack205

src/constraint_list.jl

@@ -149,6 +149,9 @@ Base.zip(cons::ConstraintList) = zip(cons.inds, cons.constraints)
 @inline Base.getindex(cons::ConstraintList, i::Int) = cons.constraints[i]
 Base.getindex(cons::ConstraintList, I) = cons.constraints[I]
 
+RD.state_dim(cons::ConstraintList, k) = cons.nx[k]
+RD.control_dim(cons::ConstraintList, k) = cons.nu[k]
+
 """
 	functionsignature(::ConstraintList, i)
 

src/problem.jl

@@ -146,7 +146,7 @@ of the trajectory, `N`.
 """
 RD.dims(prob::Problem, i::Integer) = RD.dims(prob.model[i])..., prob.N 
 
-RD.state_dim(prob::Problem, k::Integer) = state_dim(prob.model[k])
+RD.state_dim(prob::Problem, k::Integer) = state_dim(prob.model[min(k, prob.N-1)])
 RD.control_dim(prob::Problem, k::Integer) = control_dim(prob.model[k])
 
 """
@@ -213,7 +213,7 @@ Returns Vector{`RobotDynamics.DiscreteDynamics`}.
 
 Get the dynamics model at time step `k`.
 """
-@inline get_model(prob::Problem, k) = prob.model[k]
+@inline get_model(prob::Problem, k) = prob.model[min(k,prob.N-1)]
 
 "Get the objective. Returns an `AbstractObjective`."
 @inline get_objective(prob::Problem) = prob.obj