openscvx 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

openscvx/propagation.py

@@ -0,0 +1,135 @@
+import numpy as np
+
+from openscvx.config import Config
+from openscvx.integrators import solve_ivp_diffrax_prop
+
+
+def prop_aug_dy(
+    tau: float,
+    x: np.ndarray,
+    u_current: np.ndarray,
+    u_next: np.ndarray,
+    tau_init: float,
+    node: int,
+    idx_s: int,
+    state_dot: callable,
+    dis_type: str,
+    N: int,
+) -> np.ndarray:
+    x = x[None, :]
+
+    if dis_type == "ZOH":
+        beta = 0.0
+    elif dis_type == "FOH":
+        beta = (tau - tau_init) * N
+    u = u_current + beta * (u_next - u_current)
+
+    return u[:, idx_s] * state_dot(x, u[:, :-1], node).squeeze()
+
+
+def get_propagation_solver(state_dot, params):
+    def propagation_solver(V0, tau_grid, u_cur, u_next, tau_init, node, idx_s):
+        return solve_ivp_diffrax_prop(
+            f=prop_aug_dy,
+            tau_final=tau_grid[1],
+            y_0=V0,
+            args=(
+                u_cur,
+                u_next,
+                tau_init,
+                node,
+                idx_s,
+                state_dot,
+                params.dis.dis_type,
+                params.scp.n,
+            ),
+            tau_0=tau_grid[0],
+        )
+
+    return propagation_solver
+
+
+def s_to_t(u, params: Config):
+    t = [0]
+    tau = np.linspace(0, 1, params.scp.n)
+    for k in range(1, params.scp.n):
+        s_kp = u[k - 1, -1]
+        s_k = u[k, -1]
+        if params.dis.dis_type == "ZOH":
+            t.append(t[k - 1] + (tau[k] - tau[k - 1]) * (s_kp))
+        else:
+            t.append(t[k - 1] + 0.5 * (s_k + s_kp) * (tau[k] - tau[k - 1]))
+    return t
+
+
+def t_to_tau(u, t, u_nodal, t_nodal, params: Config):
+    u_lam = lambda new_t: np.array(
+        [np.interp(new_t, t_nodal, u[:, i]) for i in range(u.shape[1])]
+    ).T
+    u = np.array([u_lam(t_i) for t_i in t])
+
+    tau = np.zeros(len(t))
+    tau_nodal = np.linspace(0, 1, params.scp.n)
+    for k in range(1, len(t)):
+        k_nodal = np.where(t_nodal < t[k])[0][-1]
+        s_kp = u_nodal[k_nodal, -1]
+        tp = t_nodal[k_nodal]
+        tau_p = tau_nodal[k_nodal]
+
+        s_k = u[k, -1]
+        if params.dis.dis_type == "ZOH":
+            tau[k] = tau_p + (t[k] - tp) / s_kp
+        else:
+            tau[k] = tau_p + 2 * (t[k] - tp) / (s_k + s_kp)
+    return tau, u
+
+
+def simulate_nonlinear_time(x_0, u, tau_vals, t, params, propagation_solver):
+    states = np.empty(
+        (x_0.shape[0], 0)
+    )  # Initialize states as a 2D array with shape (n, 0)
+
+    tau = np.linspace(0, 1, params.scp.n)
+
+    u_lam = lambda new_t: np.array(
+        [np.interp(new_t, t, u[:, i]) for i in range(u.shape[1])]
+    ).T
+
+    # Bin the tau_vals into with respect to the uniform tau grid, tau
+    tau_inds = np.digitize(tau_vals, tau) - 1
+    # Force the last indice to be in the same bin as the previous ones
+    tau_inds = np.where(tau_inds == params.scp.n - 1, params.scp.n - 2, tau_inds)
+
+    prev_count = 0
+
+    for k in range(params.scp.n - 1):
+        controls_current = np.squeeze(u_lam(t[k]))[None, :]
+        controls_next = np.squeeze(u_lam(t[k + 1]))[None, :]
+
+        # Create a mask
+        mask = (tau_inds >= k) & (tau_inds < k + 1)
+
+        count = np.sum(mask)
+
+        # Use count to grab the first count number of elements
+        tau_cur = tau_vals[prev_count : prev_count + count]
+
+        sol = propagation_solver(
+            x_0,
+            (tau[k], tau[k + 1]),
+            controls_current,
+            controls_next,
+            np.array([[tau[k]]]),
+            np.array([[k]]),
+            params.sim.idx_s.stop,
+        )
+
+        x = sol.ys
+        for tau_i in tau_cur:
+            new_state = sol.evaluate(tau_i).reshape(-1, 1)  # Ensure new_state is 2D
+            states = np.concatenate([states, new_state], axis=1)
+
+        x_0 = x[-1]
+        prev_count += count
+
+    return states.T

openscvx/ptr.py

@@ -0,0 +1,149 @@
+import numpy as np
+import numpy.linalg as la
+import cvxpy as cp
+import pickle
+import time
+
+from openscvx.config import Config
+
+import warnings
+warnings.filterwarnings("ignore")
+
+def PTR_init(ocp: cp.Problem, discretization_solver: callable, params: Config):
+    if params.cvx.cvxpygen:
+        from solver.cpg_solver import cpg_solve
+        with open('solver/problem.pickle', 'rb') as f:
+            prob = pickle.load(f)
+    else:
+        cpg_solve = None
+
+    # Solve a dumb problem to intilize DPP and JAX jacobians
+    _ = PTR_subproblem(cpg_solve, params.sim.x_bar, params.sim.u_bar, discretization_solver, ocp, params)
+
+    return cpg_solve
+
+def PTR_main(params: Config, prob: cp.Problem, aug_dy: callable, cpg_solve, emitter_function) -> dict:
+    J_vb = 1E2
+    J_vc = 1E2
+    J_tr = 1E2
+
+    x_bar = params.sim.x_bar
+    u_bar = params.sim.u_bar
+
+    scp_trajs = [x_bar]
+    scp_controls = [u_bar]
+    V_multi_shoot_traj = []
+
+    k = 1
+
+    while k <= params.scp.k_max and ((J_tr >= params.scp.ep_tr) or (J_vb >= params.scp.ep_vb) or (J_vc >= params.scp.ep_vc)):
+        x, u, t, J_total, J_vb_vec, J_vc_vec, J_tr_vec, prob_stat, V_multi_shoot, subprop_time, dis_time = PTR_subproblem(cpg_solve, x_bar, u_bar, aug_dy, prob, params)
+
+        V_multi_shoot_traj.append(V_multi_shoot)
+
+        x_bar = x
+        u_bar = u
+
+        J_tr = np.sum(np.array(J_tr_vec))
+        J_vb = np.sum(np.array(J_vb_vec))
+        J_vc = np.sum(np.array(J_vc_vec))
+        scp_trajs.append(x)
+        scp_controls.append(u)
+
+        params.scp.w_tr = min(params.scp.w_tr * params.scp.w_tr_adapt, params.scp.w_tr_max)
+        if k > params.scp.cost_drop:
+            params.scp.lam_cost = params.scp.lam_cost * params.scp.cost_relax
+
+        emitter_function(
+            {
+                "iter": k,
+                "dis_time": dis_time * 1000.0,
+                "subprop_time": subprop_time * 1000.0,
+                "J_total": J_total,
+                "J_tr": J_tr,
+                "J_vb": J_vb,
+                "J_vc": J_vc,
+                "cost": t[-1],
+                "prob_stat": prob_stat,
+            }
+        )
+
+        k += 1
+
+    result = dict(
+        converged = k <= params.scp.k_max,
+        t_final = x[:,params.sim.idx_t][-1],
+        u = u,
+        x = x,
+        x_history = scp_trajs,
+        u_history = scp_controls,
+        discretization_history = V_multi_shoot_traj,
+        J_tr_history = J_tr_vec,
+        J_vb_history = J_vb_vec,
+        J_vc_history = J_vc_vec,
+    )
+    return result
+
+
+def PTR_subproblem(cpg_solve, x_bar, u_bar, aug_dy, prob, params: Config):
+    prob.param_dict['x_bar'].value = x_bar
+    prob.param_dict['u_bar'].value = u_bar
+    
+    t0 = time.time()
+    A_bar, B_bar, C_bar, z_bar, V_multi_shoot = aug_dy(x_bar, u_bar.astype(float))
+
+    prob.param_dict['A_d'].value = A_bar.__array__()
+    prob.param_dict['B_d'].value = B_bar.__array__()
+    prob.param_dict['C_d'].value = C_bar.__array__()
+    prob.param_dict['z_d'].value = z_bar.__array__()
+    dis_time = time.time() - t0
+
+    if params.sim.constraints_nodal:
+        for g_id, constraint in enumerate(params.sim.constraints_nodal):
+            if not constraint.convex:
+                prob.param_dict['g_' + str(g_id)].value = np.asarray(constraint.g(x_bar, u_bar))
+                prob.param_dict['grad_g_x_' + str(g_id)].value = np.asarray(constraint.grad_g_x(x_bar, u_bar))
+                prob.param_dict['grad_g_u_' + str(g_id)].value = np.asarray(constraint.grad_g_u(x_bar, u_bar))
+    
+    prob.param_dict['w_tr'].value = params.scp.w_tr
+    prob.param_dict['lam_cost'].value = params.scp.lam_cost
+
+    if params.cvx.cvxpygen:
+        t0 = time.time()
+        prob.register_solve('CPG', cpg_solve)
+        prob.solve(method = 'CPG', **params.cvx.solver_args)
+        subprop_time = time.time() - t0
+    else:
+        t0 = time.time()
+        prob.solve(solver = params.cvx.solver, enforce_dpp = True, **params.cvx.solver_args)
+        subprop_time = time.time() - t0
+
+    x = (params.sim.S_x @ prob.var_dict['x'].value.T + np.expand_dims(params.sim.c_x, axis = 1)).T
+    u = (params.sim.S_u @ prob.var_dict['u'].value.T + np.expand_dims(params.sim.c_u, axis = 1)).T
+
+    i = 0
+    costs = [0]
+    for type in params.sim.final_state.type:
+        if type == 'Minimize':
+            costs += x[:,i]
+        if type == 'Maximize':
+            costs -= x[:,i]
+        i += 1
+
+    # Create the block diagonal matrix using jax.numpy.block
+    inv_block_diag = np.block([
+        [params.sim.inv_S_x, np.zeros((params.sim.inv_S_x.shape[0], params.sim.inv_S_u.shape[1]))],
+        [np.zeros((params.sim.inv_S_u.shape[0], params.sim.inv_S_x.shape[1])), params.sim.inv_S_u]
+    ])
+
+    # Calculate J_tr_vec using the JAX-compatible block diagonal matrix
+    J_tr_vec = la.norm(inv_block_diag @ np.hstack((x - x_bar, u - u_bar)).T, axis=0)**2
+    J_vc_vec = np.sum(np.abs(prob.var_dict['nu'].value), axis = 1)
+    
+    id_ncvx = 0
+    J_vb_vec = 0
+    for constraint in params.sim.constraints_nodal:
+        if constraint.convex == False:
+            J_vb_vec += np.maximum(0, prob.var_dict['nu_vb_' + str(id_ncvx)].value)
+            id_ncvx += 1
+    return x, u, costs, prob.value, J_vb_vec, J_vc_vec, J_tr_vec, prob.status, V_multi_shoot, subprop_time, dis_time

openscvx/trajoptproblem.py

@@ -0,0 +1,336 @@
+import jax.numpy as jnp
+from typing import List
+import queue
+import threading
+import time
+
+import cvxpy as cp
+import jax
+import numpy as np
+
+from openscvx.config import (
+    ScpConfig,
+    SimConfig,
+    ConvexSolverConfig,
+    DiscretizationConfig,
+    PropagationConfig,
+    DevConfig,
+    Config,
+)
+from openscvx.dynamics import get_augmented_dynamics, get_jacobians
+from openscvx.constraints.violation import get_g_funcs
+from openscvx.augmentation import sort_ctcs_constraints
+from openscvx.discretization import get_discretization_solver
+from openscvx.propagation import get_propagation_solver
+from openscvx.constraints.boundary import BoundaryConstraint
+from openscvx.constraints.ctcs import CTCSConstraint
+from openscvx.constraints.nodal import NodalConstraint
+from openscvx.ptr import PTR_init, PTR_main
+from openscvx.post_processing import propagate_trajectory_results
+from openscvx.ocp import OptimalControlProblem
+from openscvx import io
+
+
+# TODO: (norrisg) Decide whether to have constraints`, `cost`, alongside `dynamics`, ` etc.
+class TrajOptProblem:
+    def __init__(
+        self,
+        dynamics: callable,
+        constraints: List[callable],
+        idx_time: int,
+        N: int,
+        time_init: float,
+        x_guess: jnp.ndarray,
+        u_guess: jnp.ndarray,
+        initial_state: BoundaryConstraint,
+        final_state: BoundaryConstraint,
+        x_max: jnp.ndarray,
+        x_min: jnp.ndarray,
+        u_max: jnp.ndarray,
+        u_min: jnp.ndarray,
+        scp: ScpConfig = None,
+        dis: DiscretizationConfig = None,
+        prp: PropagationConfig = None,
+        sim: SimConfig = None,
+        dev: DevConfig = None,
+        cvx: ConvexSolverConfig = None,
+        licq_min=0.0,
+        licq_max=1e-4,
+        time_dilation_factor_min=0.3,
+        time_dilation_factor_max=3.0,
+    ):
+
+        # TODO (norrisg) move this into some augmentation function, if we want to make this be executed after the init (i.e. within problem.initialize) need to rethink how problem is defined
+        constraints_ctcs = []
+        constraints_nodal = []
+        # TODO: (norrisg) change back to using isinstance once on PyPi
+        for constraint in constraints:
+            if type(constraint).__name__ == CTCSConstraint.__name__:
+                constraints_ctcs.append(
+                    constraint
+                )
+            elif type(constraint).__name__ == NodalConstraint.__name__:
+                constraints_nodal.append(
+                    constraint
+                )
+            else:
+                raise ValueError(
+                    f"Unknown constraint type: {type(constraint)}, All constraints must be decorated with @ctcs or @nodal"
+                )
+
+        constraints_ctcs, node_intervals, num_augmented_states = sort_ctcs_constraints(constraints_ctcs, N)
+
+        # Index tracking
+        idx_x_true = slice(0, len(x_max))
+        idx_u_true = slice(0, len(u_max))
+        idx_constraint_violation = slice(
+            idx_x_true.stop, idx_x_true.stop + num_augmented_states
+        )
+
+        idx_time_dilation = slice(idx_u_true.stop, idx_u_true.stop + 1)
+
+        # check that idx_time is in the correct range
+        assert idx_time >= 0 and idx_time < len(
+            x_max
+        ), "idx_time must be in the range of the state vector and non-negative"
+        idx_time = slice(idx_time, idx_time + 1)
+
+        x_min_augmented = np.hstack([x_min, np.repeat(licq_min, num_augmented_states)])
+        x_max_augmented = np.hstack([x_max, np.repeat(licq_max, num_augmented_states)])
+
+        u_min_augmented = np.hstack([u_min, time_dilation_factor_min * time_init])
+        u_max_augmented = np.hstack([u_max, time_dilation_factor_max * time_init])
+
+        x_bar_augmented = np.hstack([x_guess, np.full((x_guess.shape[0], num_augmented_states), 0)])
+        u_bar_augmented = np.hstack(
+            [u_guess, np.full((u_guess.shape[0], 1), time_init)]
+        )
+
+        if dis is None:
+            dis = DiscretizationConfig()
+
+        if sim is None:
+            sim = SimConfig(
+                x_bar=x_bar_augmented,
+                u_bar=u_bar_augmented,
+                initial_state=initial_state,
+                final_state=final_state,
+                max_state=x_max_augmented,
+                min_state=x_min_augmented,
+                max_control=u_max_augmented,
+                min_control=u_min_augmented,
+                total_time=time_init,
+                n_states=len(x_max),
+                idx_x_true=idx_x_true,
+                idx_u_true=idx_u_true,
+                idx_t=idx_time,
+                idx_y=idx_constraint_violation,
+                idx_s=idx_time_dilation,
+                ctcs_node_intervals=node_intervals,
+            )
+
+        if scp is None:
+            scp = ScpConfig(
+                n=N,
+                k_max=200,
+                w_tr=1e1,  # Weight on the Trust Reigon
+                lam_cost=1e1,  # Weight on the Nonlinear Cost
+                lam_vc=1e2,  # Weight on the Virtual Control Objective
+                lam_vb=0e0,  # Weight on the Virtual Buffer Objective (only for penalized nodal constraints)
+                ep_tr=1e-4,  # Trust Region Tolerance
+                ep_vb=1e-4,  # Virtual Control Tolerance
+                ep_vc=1e-8,  # Virtual Control Tolerance for CTCS
+                cost_drop=4,  # SCP iteration to relax minimal final time objective
+                cost_relax=0.5,  # Minimal Time Relaxation Factor
+                w_tr_adapt=1.2,  # Trust Region Adaptation Factor
+                w_tr_max_scaling_factor=1e2,  # Maximum Trust Region Weight
+            )
+        else:
+            assert (
+                self.scp.n == N
+            ), "Number of segments must be the same as in the config"
+
+        if dev is None:
+            dev = DevConfig()
+        if cvx is None:
+            cvx = ConvexSolverConfig()
+        if prp is None:
+            prp = PropagationConfig()
+
+        sim.constraints_ctcs = constraints_ctcs
+        sim.constraints_nodal = constraints_nodal
+
+        g_funcs = get_g_funcs(constraints_ctcs)
+        self.dynamics_augmented = get_augmented_dynamics(dynamics, g_funcs, idx_x_true, idx_u_true)
+        self.A_uncompiled, self.B_uncompiled = get_jacobians(self.dynamics_augmented)
+
+        self.params = Config(
+            sim=sim,
+            scp=scp,
+            dis=dis,
+            dev=dev,
+            cvx=cvx,
+            prp=prp,
+        )
+
+        self.optimal_control_problem: cp.Problem = None
+        self.discretization_solver: callable = None
+        self.cpg_solve = None
+
+        # set up emitter & thread only if printing is enabled
+        if self.params.dev.printing:
+            self.print_queue      = queue.Queue()
+            self.emitter_function = lambda data: self.print_queue.put(data)
+            self.print_thread     = threading.Thread(
+                target=io.intermediate,
+                args=(self.print_queue, self.params),
+                daemon=True,
+            )
+            self.print_thread.start()
+        else:
+            # no-op emitter; nothing ever gets queued or printed
+            self.emitter_function = lambda data: None
+
+
+        self.timing_init = None
+        self.timing_solve = None
+        self.timing_post = None
+
+    def initialize(self):
+        io.intro()
+
+        # Enable the profiler
+        if self.params.dev.profiling:
+            import cProfile
+
+            pr = cProfile.Profile()
+            pr.enable()
+
+        t_0_while = time.time()
+        # Ensure parameter sizes and normalization are correct
+        self.params.scp.__post_init__()
+        self.params.sim.__post_init__()
+
+        # Compile dynamics and jacobians
+        self.state_dot = jax.vmap(self.dynamics_augmented)
+        self.A = jax.jit(jax.vmap(self.A_uncompiled, in_axes=(0, 0, 0)))
+        self.B = jax.jit(jax.vmap(self.B_uncompiled, in_axes=(0, 0, 0)))
+        # TODO: (norrisg) Could consider using dataclass just to hold dynamics and jacobians
+        # TODO: (norrisg) Consider writing the compiled versions into the same variables?
+        # Otherwise if have a dataclass could have 2 instances, one for compied and one for uncompiled
+
+        # Generate solvers and optimal control problem
+        self.discretization_solver = get_discretization_solver(
+            self.state_dot, self.A, self.B, self.params
+        )
+        self.propagation_solver = get_propagation_solver(self.state_dot, self.params)
+        self.optimal_control_problem = OptimalControlProblem(self.params)
+
+        # Initialize the PTR loop
+        self.cpg_solve = PTR_init(
+            self.optimal_control_problem,
+            self.discretization_solver,
+            self.params,
+        )
+
+        # Compile the solvers
+        if not self.params.dev.debug:
+            self.discretization_solver = (
+                jax.jit(self.discretization_solver)
+                .lower(
+                    np.ones((self.params.scp.n, self.params.sim.n_states)),
+                    np.ones((self.params.scp.n, self.params.sim.n_controls)),
+                )
+                .compile()
+            )
+
+        self.propagation_solver = (
+            jax.jit(self.propagation_solver)
+            .lower(
+                np.ones((self.params.sim.n_states)),
+                (0.0, 0.0),
+                np.ones((1, self.params.sim.n_controls)),
+                np.ones((1, self.params.sim.n_controls)),
+                np.ones((1, 1)),
+                np.ones((1, 1)).astype("int"),
+                0,
+            )
+            .compile()
+        )
+
+        t_f_while = time.time()
+        self.timing_init = t_f_while - t_0_while
+        print("Total Initialization Time: ", self.timing_init)
+
+        if self.params.dev.profiling:
+            pr.disable()
+            # Save results so it can be viusualized with snakeviz
+            pr.dump_stats("profiling_initialize.prof")
+
+    def solve(self):
+        # Ensure parameter sizes and normalization are correct
+        self.params.scp.__post_init__()
+        self.params.sim.__post_init__()
+
+        if self.optimal_control_problem is None or self.discretization_solver is None:
+            raise ValueError(
+                "Problem has not been initialized. Call initialize() before solve()"
+            )
+
+        # Enable the profiler
+        if self.params.dev.profiling:
+            import cProfile
+
+            pr = cProfile.Profile()
+            pr.enable()
+
+        t_0_while = time.time()
+        # Print top header for solver results
+        io.header()
+
+        result = PTR_main(
+            self.params,
+            self.optimal_control_problem,
+            self.discretization_solver,
+            self.cpg_solve,
+            self.emitter_function,
+        )
+
+        t_f_while = time.time()
+        self.timing_solve = t_f_while - t_0_while
+
+        while self.print_queue.qsize() > 0:
+            time.sleep(0.1)
+
+        # Print bottom footer for solver results as well as total computation time
+        io.footer(self.timing_solve)
+
+        # Disable the profiler
+        if self.params.dev.profiling:
+            pr.disable()
+            # Save results so it can be viusualized with snakeviz
+            pr.dump_stats("profiling_solve.prof")
+
+        return result
+
+    def post_process(self, result):
+        # Enable the profiler
+        if self.params.dev.profiling:
+            import cProfile
+
+            pr = cProfile.Profile()
+            pr.enable()
+
+        t_0_post = time.time()
+        result = propagate_trajectory_results(self.params, result, self.propagation_solver)
+        t_f_post = time.time()
+
+        self.timing_post = t_f_post - t_0_post
+        print("Total Post Processing Time: ", self.timing_post)
+
+        # Disable the profiler
+        if self.params.dev.profiling:
+            pr.disable()
+            # Save results so it can be viusualized with snakeviz
+            pr.dump_stats("profiling_postprocess.prof")
+        return result

openscvx/utils.py

@@ -0,0 +1,80 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+def qdcm(q: jnp.ndarray) -> jnp.ndarray:
+    # Convert a quaternion to a direction cosine matrix
+    q_norm = (q[0] ** 2 + q[1] ** 2 + q[2] ** 2 + q[3] ** 2) ** 0.5
+    w, x, y, z = q / q_norm
+    return jnp.array(
+        [
+            [1 - 2 * (y**2 + z**2), 2 * (x * y - z * w), 2 * (x * z + y * w)],
+            [2 * (x * y + z * w), 1 - 2 * (x**2 + z**2), 2 * (y * z - x * w)],
+            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x**2 + y**2)],
+        ]
+    )
+
+
+def SSMP(w: jnp.ndarray):
+    # Convert an angular rate to a 4 x 4 skew symetric matrix
+    x, y, z = w
+    return jnp.array([[0, -x, -y, -z], [x, 0, z, -y], [y, -z, 0, x], [z, y, -x, 0]])
+
+
+def SSM(w: jnp.ndarray):
+    # Convert an angular rate to a 3 x 3 skew symetric matrix
+    x, y, z = w
+    return jnp.array([[0, -z, y], [z, 0, -x], [-y, x, 0]])
+
+
+def generate_orthogonal_unit_vectors(vectors=None):
+    """
+    Generates 3 orthogonal unit vectors to model the axis of the ellipsoid via QR decomposition
+
+    Parameters:
+    vectors (np.ndarray): Optional, axes of the ellipsoid to be orthonormalized.
+                            If none specified generates randomly.
+
+    Returns:
+    np.ndarray: A 3x3 matrix where each column is a unit vector.
+    """
+    if vectors is None:
+        # Create a random key
+        key = jax.random.PRNGKey(0)
+
+        # Generate a 3x3 array of random numbers uniformly distributed between 0 and 1
+        vectors = jax.random.uniform(key, (3, 3))
+    Q, _ = jnp.linalg.qr(vectors)
+    return Q
+
+
+rot = np.array(
+    [
+        [np.cos(np.pi / 2), np.sin(np.pi / 2), 0],
+        [-np.sin(np.pi / 2), np.cos(np.pi / 2), 0],
+        [0, 0, 1],
+    ]
+)
+def gen_vertices(center, radii):
+    """
+    Obtains the vertices of the gate.
+    """
+    vertices = []
+    vertices.append(center + rot @ [radii[0], 0, radii[2]])
+    vertices.append(center + rot @ [-radii[0], 0, radii[2]])
+    vertices.append(center + rot @ [-radii[0], 0, -radii[2]])
+    vertices.append(center + rot @ [radii[0], 0, -radii[2]])
+    return vertices
+
+
+# TODO (haynec): make this less hardcoded
+def get_kp_pose(t, init_pose):
+    loop_time = 40.0
+    loop_radius = 20.0
+
+    t_angle = t / loop_time * (2 * jnp.pi)
+    x = loop_radius * jnp.sin(t_angle)
+    y = x * jnp.cos(t_angle)
+    z = 0.5 * x * jnp.sin(t_angle)
+    return jnp.array([x, y, z]).T + init_pose