Lab 06-03: Trajectory Optimization

Objectives

By the end of this lab, you will be able to:

• Formulate trajectory optimization as a constrained optimization problem
• Implement direct collocation for trajectory planning
• Handle path constraints and dynamics constraints
• Optimize for multiple objectives (time, energy, smoothness)

Prerequisites

• Completed Labs 06-01 and 06-02 (path planning basics)
• Understanding of numerical optimization
• Familiarity with robot dynamics (Module 05)

Materials

Type	Name	Tier	Notes
software	Python 3.10+	required	Programming environment
software	NumPy, SciPy	required	Optimization and linear algebra
software	Matplotlib	required	Visualization
simulation	MuJoCo 3.0+	optional	For dynamics verification

Background

Trajectory Optimization Problem

Unlike path planning which finds a sequence of positions, trajectory optimization finds time-parameterized motion:

Minimize: $J = \int_0^T L(q(t), \dot{q}(t), \ddot{q}(t), u(t)) dt$

Subject to:

• Dynamics: $M(q)\ddot{q} + C(q,\dot{q})\dot{q} + g(q) = u$
• Initial/final conditions: $q(0) = q_0$, $q(T) = q_f$
• Path constraints: $c(q(t)) \leq 0$
• Control limits: $u_{min} \leq u \leq u_{max}$

Direct Collocation

Discretize trajectory into N waypoints and optimize all at once:

• Decision variables: $[q_1, q_2, ..., q_N, \dot{q}_1, ..., \dot{q}_N, u_1, ..., u_N, T]$
• Enforce dynamics at collocation points
• Use sparse optimization solvers

Instructions

Step 1: Define the Problem Structure

Set up a simple trajectory optimization problem:

import numpy as np
from scipy.optimize import minimize, NonlinearConstraint
import matplotlib.pyplot as plt

class TrajectoryOptimizer:
    """
    Direct collocation trajectory optimizer for 2-DOF arm.
    """

    def __init__(self, n_joints=2, n_waypoints=20):
        """
        Initialize trajectory optimizer.

        Args:
            n_joints: Number of robot joints
            n_waypoints: Number of discretization points
        """
        self.n_joints = n_joints
        self.n_waypoints = n_waypoints

        # Robot parameters
        self.link_masses = [2.0, 1.5]  # kg
        self.link_lengths = [0.5, 0.4]  # m
        self.link_inertias = [0.042, 0.02]  # kg*m^2

        # Limits
        self.q_min = np.array([-np.pi, -np.pi])
        self.q_max = np.array([np.pi, np.pi])
        self.qd_max = np.array([5.0, 5.0])  # rad/s
        self.u_max = np.array([50.0, 30.0])  # Nm

        # Time bounds
        self.T_min = 0.5  # seconds
        self.T_max = 5.0

    def pack_variables(self, q_traj, qd_traj, u_traj, T):
        """Pack decision variables into single vector."""
        return np.concatenate([
            q_traj.flatten(),
            qd_traj.flatten(),
            u_traj.flatten(),
            [T]
        ])

    def unpack_variables(self, x):
        """Unpack decision variables from vector."""
        n, N = self.n_joints, self.n_waypoints

        idx = 0
        q_traj = x[idx:idx + n*N].reshape(N, n)
        idx += n*N

        qd_traj = x[idx:idx + n*N].reshape(N, n)
        idx += n*N

        u_traj = x[idx:idx + n*N].reshape(N, n)
        idx += n*N

        T = x[idx]

        return q_traj, qd_traj, u_traj, T

# Create optimizer
opt = TrajectoryOptimizer(n_joints=2, n_waypoints=20)

# Test packing/unpacking
q_test = np.random.randn(20, 2)
qd_test = np.random.randn(20, 2)
u_test = np.random.randn(20, 2)
T_test = 2.0

x = opt.pack_variables(q_test, qd_test, u_test, T_test)
q_rec, qd_rec, u_rec, T_rec = opt.unpack_variables(x)

print(f"Total decision variables: {len(x)}")
print(f"Pack/unpack consistent: {np.allclose(q_test, q_rec)}")

Expected: 121 decision variables (202 + 202 + 20*2 + 1), pack/unpack returns True.

Step 2: Define Cost Function

Implement cost function for optimization:

def compute_cost(self, x, weights={'time': 1.0, 'energy': 0.1, 'smoothness': 0.01}):
    """
    Compute trajectory cost.

    Cost = w_T * T + w_E * ∫u²dt + w_S * ∫(du/dt)²dt

    Args:
        x: Decision variable vector
        weights: Dictionary of cost weights

    Returns:
        Total cost (scalar)
    """
    q_traj, qd_traj, u_traj, T = self.unpack_variables(x)
    N = self.n_waypoints
    dt = T / (N - 1)

    cost = 0.0

    # Time cost
    cost += weights['time'] * T

    # Energy cost (sum of squared torques)
    energy = np.sum(u_traj**2) * dt
    cost += weights['energy'] * energy

    # Smoothness cost (sum of squared torque derivatives)
```python
```python
    if N > 1:

        du = np.diff(u_traj, axis=0) / dt
        smoothness = np.sum(du**2) * dt
        cost += weights['smoothness'] * smoothness

    return cost

# Add method to class
TrajectoryOptimizer.compute_cost = compute_cost

# Test cost computation
q_init = np.linspace([0, 0], [1.5, 1.0], 20)
qd_init = np.zeros((20, 2))
u_init = np.zeros((20, 2))
T_init = 2.0

x_init = opt.pack_variables(q_init, qd_init, u_init, T_init)
cost = opt.compute_cost(x_init)
print(f"Initial cost: {cost:.4f}")

Expected: Cost value computed without errors.

Step 3: Implement Dynamics Constraints

Enforce robot dynamics using collocation:

def compute_dynamics_residual(self, q, qd, qdd, u):
    """
    Compute dynamics residual: M*qdd + C*qd + g - u

    For 2-DOF arm with simplified dynamics.
    """
    m1, m2 = self.link_masses
    l1, l2 = self.link_lengths
    lc1, lc2 = l1/2, l2/2
    I1, I2 = self.link_inertias
    g = 9.81

    q1, q2 = q
    qd1, qd2 = qd

    # Mass matrix elements
    M11 = m1*lc1**2 + I1 + m2*(l1**2 + lc2**2 + 2*l1*lc2*np.cos(q2)) + I2
    M12 = m2*(lc2**2 + l1*lc2*np.cos(q2)) + I2
    M21 = M12
    M22 = m2*lc2**2 + I2

    M = np.array([[M11, M12], [M21, M22]])

    # Coriolis terms
    h = -m2*l1*lc2*np.sin(q2)
    C = np.array([
        [h*qd2, h*(qd1 + qd2)],
        [-h*qd1, 0]
    ])

    # Gravity
    g_vec = np.array([
        (m1*lc1 + m2*l1)*g*np.cos(q1) + m2*lc2*g*np.cos(q1 + q2),
        m2*lc2*g*np.cos(q1 + q2)
    ])

    # Dynamics residual
    residual = M @ qdd + C @ qd + g_vec - u

    return residual

def dynamics_constraints(self, x):
    """
    Compute all dynamics constraint violations.

    Uses trapezoidal collocation:
    q[k+1] = q[k] + dt/2 * (qd[k] + qd[k+1])
    qd[k+1] = qd[k] + dt/2 * (qdd[k] + qdd[k+1])

    Returns:
        Array of constraint violations (should be zero)
    """
    q_traj, qd_traj, u_traj, T = self.unpack_variables(x)
    N = self.n_waypoints
    dt = T / (N - 1)

    constraints = []

    for k in range(N - 1):
        q_k, q_kp1 = q_traj[k], q_traj[k+1]
        qd_k, qd_kp1 = qd_traj[k], qd_traj[k+1]
        u_k, u_kp1 = u_traj[k], u_traj[k+1]

        # Estimate accelerations from dynamics
        # M*qdd = u - C*qd - g
        qdd_k = self._compute_acceleration(q_k, qd_k, u_k)
        qdd_kp1 = self._compute_acceleration(q_kp1, qd_kp1, u_kp1)

        # Trapezoidal integration constraints
        # Position: q[k+1] - q[k] - dt/2*(qd[k] + qd[k+1]) = 0
        pos_defect = q_kp1 - q_k - dt/2 * (qd_k + qd_kp1)

        # Velocity: qd[k+1] - qd[k] - dt/2*(qdd[k] + qdd[k+1]) = 0
        vel_defect = qd_kp1 - qd_k - dt/2 * (qdd_k + qdd_kp1)

        constraints.extend(pos_defect.tolist())
        constraints.extend(vel_defect.tolist())

    return np.array(constraints)

def _compute_acceleration(self, q, qd, u):
    """Compute acceleration from inverse dynamics."""
    m1, m2 = self.link_masses
    l1, l2 = self.link_lengths
    lc1, lc2 = l1/2, l2/2
    I1, I2 = self.link_inertias
    g = 9.81

    q1, q2 = q
    qd1, qd2 = qd

    # Mass matrix
    M11 = m1*lc1**2 + I1 + m2*(l1**2 + lc2**2 + 2*l1*lc2*np.cos(q2)) + I2
    M12 = m2*(lc2**2 + l1*lc2*np.cos(q2)) + I2
    M22 = m2*lc2**2 + I2
    M = np.array([[M11, M12], [M12, M22]])

    # Coriolis
    h = -m2*l1*lc2*np.sin(q2)
    C_qd = np.array([
        h*qd2*(2*qd1 + qd2),
        -h*qd1*qd1
    ])

    # Gravity
    g_vec = np.array([
        (m1*lc1 + m2*l1)*g*np.cos(q1) + m2*lc2*g*np.cos(q1 + q2),
        m2*lc2*g*np.cos(q1 + q2)
    ])

    # Solve for acceleration
    qdd = np.linalg.solve(M, u - C_qd - g_vec)

    return qdd

# Add methods to class
TrajectoryOptimizer.compute_dynamics_residual = compute_dynamics_residual
TrajectoryOptimizer.dynamics_constraints = dynamics_constraints
TrajectoryOptimizer._compute_acceleration = _compute_acceleration

# Test dynamics constraints
constraints = opt.dynamics_constraints(x_init)
print(f"Number of dynamics constraints: {len(constraints)}")
print(f"Initial constraint violation (should be large): {np.linalg.norm(constraints):.4f}")

Expected: 76 dynamics constraints (19 intervals × 4 defects), large initial violation.

Step 4: Boundary and Path Constraints

Add start/end conditions and limits:

def boundary_constraints(self, x, q_start, q_goal):
    """
    Boundary condition constraints.

    q(0) = q_start, q(T) = q_goal
    qd(0) = 0, qd(T) = 0 (start and end at rest)
    """
    q_traj, qd_traj, u_traj, T = self.unpack_variables(x)

    constraints = []

    # Initial position and velocity
    constraints.extend((q_traj[0] - q_start).tolist())
    constraints.extend(qd_traj[0].tolist())

    # Final position and velocity
    constraints.extend((q_traj[-1] - q_goal).tolist())
    constraints.extend(qd_traj[-1].tolist())

    return np.array(constraints)

def get_bounds(self, q_start, q_goal):
    """
    Get variable bounds for optimization.

    Returns:
        List of (lower, upper) bounds for each variable
    """
    N = self.n_waypoints
    n = self.n_joints

    bounds = []

    # Position bounds
    for _ in range(N):
        for j in range(n):
            bounds.append((self.q_min[j], self.q_max[j]))

    # Velocity bounds
    for _ in range(N):
        for j in range(n):
            bounds.append((-self.qd_max[j], self.qd_max[j]))

    # Control bounds
    for _ in range(N):
        for j in range(n):
            bounds.append((-self.u_max[j], self.u_max[j]))

    # Time bounds
    bounds.append((self.T_min, self.T_max))

    return bounds

TrajectoryOptimizer.boundary_constraints = boundary_constraints
TrajectoryOptimizer.get_bounds = get_bounds

Expected: Boundary constraint function returns 8 constraints (4 initial + 4 final).

Step 5: Solve the Optimization

Put it all together and solve:

def optimize(self, q_start, q_goal, T_init=2.0, verbose=True):
    """
    Optimize trajectory from q_start to q_goal.

    Args:
        q_start: Initial configuration
        q_goal: Goal configuration
        T_init: Initial guess for duration
        verbose: Print progress

    Returns:
        Optimized trajectory dictionary
    """
    N = self.n_waypoints
    n = self.n_joints

    # Initialize with linear interpolation
    q_init = np.linspace(q_start, q_goal, N)
    qd_init = np.zeros((N, n))

    # Initialize controls to counteract gravity at each point
    u_init = np.zeros((N, n))
    for i in range(N):
        # Simple gravity compensation
        u_init[i] = self._gravity_compensation(q_init[i])

    x0 = self.pack_variables(q_init, qd_init, u_init, T_init)

    # Cost function
    def cost_fn(x):
        return self.compute_cost(x)

    # Constraints
    def dynamics_eq(x):
        return self.dynamics_constraints(x)

    def boundary_eq(x):
        return self.boundary_constraints(x, q_start, q_goal)

    dynamics_constraint = NonlinearConstraint(dynamics_eq, 0, 0)
    boundary_constraint = NonlinearConstraint(boundary_eq, 0, 0)

    # Bounds
    bounds = self.get_bounds(q_start, q_goal)

    # Solve
    if verbose:
        print("Starting optimization...")
        print(f"  Variables: {len(x0)}")
        print(f"  Dynamics constraints: {len(dynamics_eq(x0))}")
        print(f"  Boundary constraints: {len(boundary_eq(x0))}")

    result = minimize(
        cost_fn,
        x0,
        method='SLSQP',
        bounds=bounds,
        constraints=[
            {'type': 'eq', 'fun': dynamics_eq},
            {'type': 'eq', 'fun': boundary_eq}
        ],
        options={'maxiter': 500, 'disp': verbose}
    )

    if verbose:
        print(f"\nOptimization {'succeeded' if result.success else 'failed'}")
        print(f"Final cost: {result.fun:.4f}")

    # Unpack result
    q_opt, qd_opt, u_opt, T_opt = self.unpack_variables(result.x)

    return {
        'q': q_opt,
        'qd': qd_opt,
        'u': u_opt,
        'T': T_opt,
        'success': result.success,
        'cost': result.fun,
        'time_vector': np.linspace(0, T_opt, N)
    }

def _gravity_compensation(self, q):
    """Compute gravity compensation torque."""
    m1, m2 = self.link_masses
    l1 = self.link_lengths[0]
    lc1, lc2 = l1/2, self.link_lengths[1]/2
    g = 9.81

    q1, q2 = q

    tau1 = (m1*lc1 + m2*l1)*g*np.cos(q1) + m2*lc2*g*np.cos(q1 + q2)
    tau2 = m2*lc2*g*np.cos(q1 + q2)

    return np.array([tau1, tau2])

TrajectoryOptimizer.optimize = optimize
TrajectoryOptimizer._gravity_compensation = _gravity_compensation

# Run optimization
q_start = np.array([0.0, 0.0])
q_goal = np.array([np.pi/2, -np.pi/4])

result = opt.optimize(q_start, q_goal, T_init=2.0)

print(f"\nOptimized duration: {result['T']:.3f} seconds")
print(f"Final position error: {np.linalg.norm(result['q'][-1] - q_goal):.6f}")

Expected: Optimization converges, final position error near zero.

Step 6: Visualize Results

Create comprehensive visualization:

def visualize_trajectory(result, opt):
    """Visualize optimized trajectory."""
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))

    t = result['time_vector']
    q = result['q']
    qd = result['qd']
    u = result['u']

    # Joint positions
    axes[0, 0].plot(t, np.degrees(q[:, 0]), 'b-', linewidth=2, label='Joint 1')
    axes[0, 0].plot(t, np.degrees(q[:, 1]), 'r-', linewidth=2, label='Joint 2')
    axes[0, 0].set_xlabel('Time (s)')
    axes[0, 0].set_ylabel('Position (deg)')
    axes[0, 0].set_title('Joint Positions')
    axes[0, 0].legend()
    axes[0, 0].grid(True)

    # Joint velocities
    axes[0, 1].plot(t, np.degrees(qd[:, 0]), 'b-', linewidth=2, label='Joint 1')
    axes[0, 1].plot(t, np.degrees(qd[:, 1]), 'r-', linewidth=2, label='Joint 2')
    axes[0, 1].axhline(y=np.degrees(opt.qd_max[0]), color='k', linestyle='--', alpha=0.5)
    axes[0, 1].axhline(y=-np.degrees(opt.qd_max[0]), color='k', linestyle='--', alpha=0.5)
    axes[0, 1].set_xlabel('Time (s)')
    axes[0, 1].set_ylabel('Velocity (deg/s)')
    axes[0, 1].set_title('Joint Velocities')
    axes[0, 1].legend()
    axes[0, 1].grid(True)

    # Control torques
    axes[1, 0].plot(t, u[:, 0], 'b-', linewidth=2, label='Joint 1')
    axes[1, 0].plot(t, u[:, 1], 'r-', linewidth=2, label='Joint 2')
    axes[1, 0].axhline(y=opt.u_max[0], color='b', linestyle='--', alpha=0.5)
    axes[1, 0].axhline(y=-opt.u_max[0], color='b', linestyle='--', alpha=0.5)
    axes[1, 0].axhline(y=opt.u_max[1], color='r', linestyle='--', alpha=0.5)
    axes[1, 0].axhline(y=-opt.u_max[1], color='r', linestyle='--', alpha=0.5)
    axes[1, 0].set_xlabel('Time (s)')
    axes[1, 0].set_ylabel('Torque (Nm)')
    axes[1, 0].set_title('Control Torques')
    axes[1, 0].legend()
    axes[1, 0].grid(True)

    # Workspace path
    ax = axes[1, 1]
    l1, l2 = opt.link_lengths

    for i in range(len(q)):
        alpha = 0.2 + 0.8 * i / len(q)
        x1 = l1 * np.cos(q[i, 0])
        y1 = l1 * np.sin(q[i, 0])
        x2 = x1 + l2 * np.cos(q[i, 0] + q[i, 1])
        y2 = y1 + l2 * np.sin(q[i, 0] + q[i, 1])

        ax.plot([0, x1, x2], [0, y1, y2], 'b-', alpha=alpha, linewidth=1)

    # Start and end
    x1_s = l1 * np.cos(q[0, 0])
    y1_s = l1 * np.sin(q[0, 0])
    x2_s = x1_s + l2 * np.cos(q[0, 0] + q[0, 1])
    y2_s = y1_s + l2 * np.sin(q[0, 0] + q[0, 1])
    ax.plot([0, x1_s, x2_s], [0, y1_s, y2_s], 'g-', linewidth=3, label='Start')

    x1_e = l1 * np.cos(q[-1, 0])
    y1_e = l1 * np.sin(q[-1, 0])
    x2_e = x1_e + l2 * np.cos(q[-1, 0] + q[-1, 1])
    y2_e = y1_e + l2 * np.sin(q[-1, 0] + q[-1, 1])
    ax.plot([0, x1_e, x2_e], [0, y1_e, y2_e], 'r-', linewidth=3, label='End')

    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_title('Workspace Trajectory')
    ax.set_aspect('equal')
    ax.legend()
    ax.grid(True)

    plt.tight_layout()
    plt.savefig('trajectory_optimization.png', dpi=150)
    plt.show()

visualize_trajectory(result, opt)

Expected: Four-panel plot showing smooth position, velocity, torque profiles and workspace motion.

Expected Outcomes

After completing this lab, you should have:

1. Code artifacts:

   • trajectory_optimizer.py: Complete direct collocation optimizer
   • arm_dynamics.py: 2-DOF arm dynamics functions

2. Visual outputs:

   • trajectory_optimization.png: Multi-panel trajectory visualization

3. Understanding:

   • Trajectory optimization formulation
   • Direct collocation method
   • Trade-offs between time, energy, and smoothness

Rubric

Criterion	Points	Description
Problem Formulation	20	Correct variable packing and cost function
Dynamics Constraints	30	Proper collocation constraints implementation
Boundary Conditions	15	Start/end position and velocity constraints
Optimization	20	Successful solution finding
Visualization	15	Clear, informative trajectory plots
Total	100

Common Errors

Error: Optimization fails to converge Solution: Improve initial guess, reduce problem size, check constraint scaling.

Error: Dynamics constraints violated Solution: Verify dynamics equations, check M matrix conditioning.

Error: Solution violates bounds Solution: Ensure bounds are properly formatted, check for infeasible problem.

Extensions

For advanced students:

1. Obstacle Avoidance: Add signed distance constraints
2. Multiple Phases: Handle contact/non-contact phases
3. Warm Starting: Use previous solution to speed up replanning
4. Shooting Methods: Implement single/multiple shooting alternatives

Related Content

• Theory: Module 06 theory.md, Section 6.5 (Trajectory Optimization)
• Previous Labs: Lab 06-01 (RRT), Lab 06-02 (A*)
• Application: See Figure AI case study for real-world trajectory optimization