"""Original AA203 teaching experiments, distinct from Stanford assignments. Requires NumPy. Run `python numerical_lab.py` for independent numerical checks. All public functions are pure: they return arrays/scalars and do not control hardware, fetch data, or write files. Units are supplied by the calling model. """ import json import math import numpy as np def _finite(value, name): value = np.asarray(value, dtype=float) if not np.all(np.isfinite(value)): raise ValueError(f'{name} must contain finite values') return value def _natural(value, name): if isinstance(value, bool) or int(value) != value or value < 0: raise ValueError(f'{name} must be a nonnegative integer') return int(value) def minimum_energy_transfer(distance, duration, time): """Rest-to-rest double integrator: position, velocity, acceleration. Minimize integral u(t)^2 dt with x(0)=v(0)=v(T)=0 and x(T)=distance. Duration is strictly positive, evaluation times belong to [0, duration], and there are no acceleration or velocity constraints. Distance may be negative. Returned arrays have the same shape as `time`. """ distance = float(_finite(distance, 'distance')) duration = float(_finite(duration, 'duration')) time = _finite(time, 'time') if duration <= 0 or np.any(time < 0) or np.any(time > duration): raise ValueError('Require duration > 0 and 0 <= time <= duration') tau = time / duration x = distance * (3 * tau ** 2 - 2 * tau ** 3) v = 6 * distance / duration * tau * (1 - tau) u = 6 * distance / duration ** 2 * (1 - 2 * tau) return x, v, u def two_pulse_transfer(distance, duration, penalty=None, limit=None): """Rest-to-rest transfer using equal-duration opposite accelerations. distance [m], duration [s], and optional acceleration limit [m/s^2]. With penalty=None enforce final position exactly; otherwise minimize integral u^2 dt + penalty*(p_final-distance)^2, penalty in s^-3. Return the acceleration, final position, and squared-acceleration effort. This optimizes within a two-pulse family, not over all control histories. """ distance = float(_finite(distance, 'distance')) duration = float(_finite(duration, 'duration')) if duration <= 0: raise ValueError('duration must be positive') half = duration / 2 if limit is not None: limit = float(_finite(limit, 'limit')) if limit <= 0: raise ValueError('limit must be positive') if penalty is None: acceleration = distance / half ** 2 if limit is not None and abs(acceleration) > limit + 1e-12: raise ValueError('Requested rest-to-rest transfer is infeasible') else: penalty = float(_finite(penalty, 'penalty')) if penalty < 0: raise ValueError('penalty must be nonnegative') acceleration = penalty * half * distance / (penalty * half ** 3 + 2) if limit is not None: acceleration = float(np.clip(acceleration, -limit, limit)) final_position = acceleration * half ** 2 effort = duration * acceleration ** 2 return acceleration, final_position, effort def pulse_equality_solution(distance=1, penalty=2, final_velocity=0): """Solve the normalized two-input quadratic using its KKT linear system. Two intervals each have unit duration. Minimize u0^2+u1^2 plus penalty*(1.5*u0+0.5*u1-distance)^2, with u0+u1=final_velocity. Inputs/outputs use normalized units. No actuator bounds are imposed. Return the unique input pair and equality multiplier for L=F+lambda*h. """ distance = float(_finite(distance, 'distance')) penalty = float(_finite(penalty, 'penalty')) final_velocity = float(_finite(final_velocity, 'final_velocity')) if penalty < 0: raise ValueError('penalty must be nonnegative') position_row = np.array([1.5, 0.5]) equality = np.ones((1, 2)) hessian = 2 * np.eye(2) + 2 * penalty * np.outer( position_row, position_row) matrix = np.block([[hessian, equality.T], [equality, np.zeros((1, 1))]]) rhs = np.r_[2 * penalty * distance * position_row, final_velocity] solution = np.linalg.solve(matrix, rhs) return solution[:2], float(solution[2]) def pulse_armijo_step(acceleration, penalty=2, distance=1, sufficient_decrease=0.1, shrink=0.5): """One backtracked negative-gradient step for a normalized pulse cost. Objective is F(a)=2*a^2+penalty*(a-distance)^2. Start the step at one. Return updated a, accepted step size, and number of evaluated trials. This is one step, not a general convergence certificate. """ acceleration = float(_finite(acceleration, 'acceleration')) penalty = float(_finite(penalty, 'penalty')) distance = float(_finite(distance, 'distance')) if penalty < 0 or not 0 < sufficient_decrease < 1 or not 0 < shrink < 1: raise ValueError('Require penalty >= 0 and line-search factors in (0,1)') gradient = 4 * acceleration + 2 * penalty * (acceleration - distance) if gradient == 0: return acceleration, 0., 0 initial_cost = 2 * acceleration ** 2 + penalty * ( acceleration - distance) ** 2 step = 1. for trial in range(1, 101): candidate = acceleration - step * gradient cost = 2 * candidate ** 2 + penalty * (candidate - distance) ** 2 target = initial_cost - sufficient_decrease * step * gradient ** 2 if cost <= target: return candidate, step, trial step *= shrink raise RuntimeError('Backtracking exhausted; no descent step returned') def finite_lqr(a, b, q, r, terminal, horizon): """Finite-horizon unconstrained discrete LQR for x_next=A x+B u. Minimize sum x'Qx + u'Ru plus x_N' terminal x_N. Returns feedback gains K (u=-Kx) and value matrices P. Q/terminal must be positive semidefinite and R strictly positive definite. Zero horizon is supported. Uses linear solves, never explicit inverses. No controllability claim is made: a finite-horizon solution exists even for an uncontrollable plant. """ a, b, q, r, terminal = [_finite(v, n) for v, n in zip( (a, b, q, r, terminal), ('A', 'B', 'Q', 'R', 'terminal'))] horizon = _natural(horizon, 'horizon') if a.ndim != 2 or a.shape[0] != a.shape[1] or b.ndim != 2: raise ValueError('A must be square and B a matrix') n, m = b.shape if a.shape != (n, n) or q.shape != (n, n): raise ValueError('Incompatible state dimensions') if terminal.shape != (n, n) or r.shape != (m, m): raise ValueError('Incompatible cost dimensions') for name, cost in [('Q', q), ('R', r), ('terminal', terminal)]: if not np.allclose(cost, cost.T): raise ValueError(f'{name} must be symmetric') if np.linalg.eigvalsh(cost).min() < -1e-12: raise ValueError(f'{name} must be positive semidefinite') if np.linalg.eigvalsh(r).min() <= 0: raise ValueError('R must be positive definite') gains = np.zeros((horizon, m, n)) values = np.zeros((horizon + 1, n, n)) values[-1] = terminal for k in range(horizon - 1, -1, -1): p = values[k + 1] h = r + b.T @ p @ b gains[k] = np.linalg.solve(h, b.T @ p @ a) values[k] = q + a.T @ p @ a - a.T @ p @ b @ gains[k] values[k] = (values[k] + values[k].T) / 2 return gains, values def rollout(a, b, gains, initial, disturbances=None): """Apply supplied linear gains; optionally add one state shock per step.""" initial = _finite(initial, 'initial') gains = _finite(gains, 'gains') n = initial.size states = np.zeros((len(gains) + 1, n)) actions = np.zeros((len(gains), b.shape[1])) shocks = (np.zeros((len(gains), n)) if disturbances is None else _finite(disturbances, 'disturbances')) if shocks.shape != (len(gains), n): raise ValueError('One state disturbance vector is required per step') states[0] = initial for k, gain in enumerate(gains): actions[k] = -gain @ states[k] states[k + 1] = a @ states[k] + b @ actions[k] + shocks[k] return states, actions def nonlinear_cart_ilqr(initial=(1., 0.), horizon=30, step=.1, drag=.2, max_iterations=100, tolerance=1e-7): """Original unconstrained iLQR example for a normalized nonlinear cart. The discrete model is p_next=p+h*v, v_next=v+h*(u-drag*v^3). This Euler model is the optimization target, not an exact continuous map. Cost: half sum x' diag(1,.1) x + .05*u^2 plus half x_N' diag(20,2) x_N. Return states, controls, accepted costs and final gradient infinity norm. Full nonlinear rollout, feedback line search, and positive scalar control curvature are checked. Failure to converge raises rather than certifies. """ initial = _finite(initial, 'initial') horizon = _natural(horizon, 'horizon') max_iterations = _natural(max_iterations, 'max_iterations') step, drag = float(_finite(step, 'step')), float(_finite(drag, 'drag')) if (initial.shape != (2,) or horizon < 1 or max_iterations < 1 or step <= 0 or drag < 0 or tolerance <= 0): raise ValueError('Invalid initial state or algorithm parameters') q, terminal, r = np.diag([1., .1]), np.diag([20., 2.]), .05 b = np.array([0., step]) def dynamics(x, u): return x + step * np.array([x[1], u - drag * x[1] ** 3]) def jacobian(x): return np.array([[1., step], [0., 1. - 3*step*drag*x[1]**2]]) def cost(xs, us): return .5*(np.einsum('ki,ij,kj->', xs[:-1], q, xs[:-1]) + r*(us @ us) + xs[-1] @ terminal @ xs[-1]) def rollout_local(us, nominal=None, feed=None, gains=None, alpha=1.): xs, out = np.zeros((horizon+1, 2)), np.zeros(horizon) xs[0] = initial with np.errstate(over='ignore', invalid='ignore'): for k in range(horizon): out[k] = us[k] if nominal is None else ( us[k] + alpha*feed[k] + gains[k] @ (xs[k]-nominal[k])) xs[k+1] = dynamics(xs[k], out[k]) score = cost(xs, out) return xs, out, score xs, us, score = rollout_local(np.zeros(horizon)) history = [float(score)] for _ in range(max_iterations): vx, vxx = terminal @ xs[-1], terminal.copy() feed, gains = np.zeros(horizon), np.zeros((horizon, 2)) for k in range(horizon-1, -1, -1): a = jacobian(xs[k]) qx, qu = q @ xs[k] + a.T @ vx, r*us[k] + b @ vx qxx, quu = q + a.T @ vxx @ a, r + b @ vxx @ b qux = b @ vxx @ a if not np.isfinite(quu) or quu <= 0: raise RuntimeError('Nonpositive local control curvature') feed[k], gains[k] = -qu/quu, -qux/quu vx = qx - qux*qu/quu vxx = qxx - np.outer(qux, qux)/quu vxx = (vxx + vxx.T)/2 # Independently compute exact gradient of the supplied discrete model. adjoint, gradient = terminal @ xs[-1], np.zeros(horizon) for k in range(horizon-1, -1, -1): gradient[k] = r*us[k] + b @ adjoint adjoint = q @ xs[k] + jacobian(xs[k]).T @ adjoint norm = float(np.max(np.abs(gradient))) if norm <= tolerance: return xs, us, history, norm for alpha in 2.**(-np.arange(20)): trial_x, trial_u, trial = rollout_local(us, xs, feed, gains, alpha) if np.isfinite(trial) and trial < score: xs, us, score = trial_x, trial_u, trial history.append(float(score)) break else: raise RuntimeError('Line search failed before stationarity') raise RuntimeError('Iteration limit; no converged result returned') def discounted_value_iteration(transitions, costs, discount, tolerance=1e-10): """Solve a finite cost-minimizing MDP, with an a posteriori error bound. transitions[s,a,s_next] are probabilities; costs[s,a] are expected costs. Return value, greedy policy, and residual/(1-discount). Discount is in [0,1). The error bound concerns this supplied finite model only. """ transitions = _finite(transitions, 'transitions') costs = _finite(costs, 'costs') if transitions.ndim != 3 or costs.shape != transitions.shape[:2]: raise ValueError('Require transitions[S,A,S] and costs[S,A]') if transitions.shape[0] != transitions.shape[2]: raise ValueError('Next-state count must equal state count') if np.any(transitions < 0) or not np.allclose( transitions.sum(axis=2), 1): raise ValueError('Transition rows must be probability vectors') if not 0 <= discount < 1 or not 0 < tolerance < math.inf: raise ValueError('Require discount in [0,1) and positive tolerance') values = np.zeros(costs.shape[0]) for _ in range(100000): candidates = costs + discount * np.einsum( 'sak,k->sa', transitions, values) updated = candidates.min(axis=1) residual = np.max(np.abs(updated - values)) if residual <= tolerance * (1 - discount): return values, candidates.argmin(axis=1), residual / (1 - discount) values = updated raise RuntimeError('Iteration limit reached; no convergence claimed') def robust_predecessor(states, actions, disturbances, target, dynamics): """Exact finite-set predecessor: exists action, for every disturbance. States/actions/disturbances are finite iterables; target is a set. Equality is exact, making this suitable for integer toy systems, not discretized floating-point reachability certification for a continuous robot. """ actions, disturbances = tuple(actions), tuple(disturbances) if not actions or not disturbances: raise ValueError('Action and disturbance sets must be nonempty') return {x for x in states if any(all(dynamics(x, u, w) in target for w in disturbances) for u in actions)} def identify_linear(states, actions, next_states): """Fit x_next=A x+B u by unregularized least squares. Rows are independent transition samples. Reject rank-deficient designs because the separate A/B parameters are then unidentifiable from the data. A low residual alone does not certify extrapolation or closed-loop safety. """ states, actions, next_states = [_finite(v, n) for v, n in zip( (states, actions, next_states), ('states', 'actions', 'next_states'))] if any(v.ndim != 2 for v in (states, actions, next_states)): raise ValueError('Samples must be row matrices') if len(states) != len(actions) or states.shape != next_states.shape: raise ValueError('Incompatible sample arrays') design = np.column_stack((states, actions)) if np.linalg.matrix_rank(design) < design.shape[1]: raise ValueError('Unidentifiable parameters: design lacks full rank') fit, _, _, _ = np.linalg.lstsq(design, next_states, rcond=None) n = states.shape[1] residual = np.linalg.norm(design @ fit - next_states) return fit[:n].T, fit[n:].T, float(residual) def bernoulli_policy_gradient(logit, reward_zero, reward_one): """Exact one-step score-function gradient of expected reward. Policy chooses action one with sigmoid(logit), zero otherwise. Returns mean reward and its derivative with respect to logit; no sampled estimator is used. This is a local identity, not a multistep RL benchmark. """ logit = float(_finite(logit, 'logit')) rewards = _finite([reward_zero, reward_one], 'rewards') p = (1 / (1 + math.exp(-logit)) if logit >= 0 else math.exp(logit) / (1 + math.exp(logit))) mean = (1 - p) * rewards[0] + p * rewards[1] gradient = p * (1 - p) * (rewards[1] - rewards[0]) return float(mean), float(gradient) def bounded_energy_profile(distance, duration, limit, time): """Exact bounded rest-to-rest acceleration, squared effort, and ramp width. Returns the optimal profile for the ideal double integrator, not a feedback safety controller. A zero-width ramp at minimum duration is a limiting bang-bang case; zero displacement is handled separately. """ distance, duration, limit = _finite([distance, duration, limit], 'parameters') time = _finite(time, 'time') if duration <= 0 or limit <= 0 or np.any(time < 0) or np.any(time > duration): raise ValueError('Require positive duration/limit and times within horizon') magnitude = abs(distance) if magnitude == 0: return np.zeros_like(time), 0., duration / 2 minimum_time = 2 * np.sqrt(magnitude / limit) if duration < minimum_time: raise ValueError('The requested bounded transfer is infeasible') if duration >= np.sqrt(6 * magnitude / limit): control = 6 * distance / duration ** 2 * (1 - 2 * time / duration) return control, float(12 * distance ** 2 / duration ** 3), duration / 2 width = np.sqrt(max(0., 3 * (duration ** 2 / 4 - magnitude / limit))) if width == 0: control = np.where(time < duration / 2, limit, -limit) * np.sign(distance) else: control = np.clip(limit / width * (duration / 2 - time), -limit, limit) control *= np.sign(distance) effort = limit ** 2 * duration - 4 * limit ** 2 * width / 3 return control, float(effort), float(width) def free_time_transfer(distance, time_weight=1., effort_weight=1.): """Return optimal positive duration and cost for an unbounded rest transfer. Objective is time_weight*T**2/2 + effort_weight*integral(u**2)/2. Requires nonzero signed displacement and strictly positive weights. No actuator bound, control penalty discontinuity, or friction is modeled. """ distance, time_weight, effort_weight = _finite( [distance, time_weight, effort_weight], 'free-time parameters') if distance == 0 or time_weight <= 0 or effort_weight <= 0: raise ValueError('Require nonzero displacement and positive weights') duration = (18 * effort_weight * distance ** 2 / time_weight) ** 0.2 cost = time_weight * duration ** 2 / 2 + 6 * effort_weight * distance ** 2 / duration ** 3 return float(duration), float(cost) def run_checks(): """Check derivations against independent optimizations and boundaries.""" # Original lecture-1 example: hard endpoints versus a soft position cost. hard = two_pulse_transfer(1, 2, limit=1) soft = two_pulse_transfer(1, 2, penalty=2, limit=1) assert np.allclose(hard, [1, 1, 2]) assert np.allclose(soft, [0.5, 0.5, 0.5]) grid = np.linspace(-1, 1, 20001) grid_cost = 2 * grid ** 2 + 2 * (grid - 1) ** 2 assert abs(grid[np.argmin(grid_cost)] - soft[0]) < 1e-12 assert np.allclose(two_pulse_transfer(-1, 2, limit=1), [-1, -1, 2]) assert np.allclose(two_pulse_transfer(1, 2, penalty=0), [0, 0, 0]) assert np.allclose(two_pulse_transfer(1, 2, penalty=2, limit=0.2), [0.2, 0.2, 0.08]) try: two_pulse_transfer(1, 1.9, limit=1) except ValueError: pass else: raise AssertionError('Infeasible hard transfer was accepted') # A matrix KKT solve independently reproduces the scalar elimination. equality_actions, equality_multiplier = pulse_equality_solution() assert np.allclose(equality_actions, [0.5, -0.5]) assert abs(equality_multiplier - 2) < 1e-12 for target in (-2, 0, 1.5): actions, _ = pulse_equality_solution(target, penalty=3) pulse = two_pulse_transfer(target, 2, penalty=3)[0] assert np.allclose(actions, [pulse, -pulse]) assert np.allclose(pulse_equality_solution(penalty=0)[0], 0) step_result = pulse_armijo_step(0) assert np.allclose(step_result, [0.5, 0.125, 4]) assert pulse_armijo_step(0.5) == (0.5, 0., 0) sensitivity_eps = 1e-5 costs = [] for velocity in (-sensitivity_eps, sensitivity_eps): actions, _ = pulse_equality_solution(final_velocity=velocity) final_position = np.array([1.5, 0.5]) @ actions costs.append(actions @ actions + 2 * (final_position - 1) ** 2) sensitivity = (costs[1] - costs[0]) / (2 * sensitivity_eps) assert abs(sensitivity + equality_multiplier) < 1e-9 # Independent integration confirms the analytical minimum-energy cost. time = np.linspace(0, 2, 10001) x, v, u = minimum_energy_transfer(1, 2, time) assert np.allclose([x[0], v[0], x[-1], v[-1]], [0, 0, 1, 0]) energy = np.trapezoid(u ** 2, time) assert abs(energy - 12 / 2 ** 3) < 1e-7 assert np.allclose(minimum_energy_transfer(0, 2, time), 0) # Lecture 3: distinguish a full increment from its first variation. tau = np.linspace(0, 1, 20001) eta, eta_dot = tau * (1 - tau), 1 - 2 * tau first_variation = np.trapezoid(tau * eta + eta_dot, tau) assert abs(first_variation - 1 / 12) < 1e-9 base_cost = np.trapezoid((tau ** 2 + 1) / 2, tau) for perturbation in (-0.2, 0.01, 0.3): perturbed_cost = np.trapezoid( ((tau + perturbation * eta) ** 2 + (1 + perturbation * eta_dot) ** 2) / 2, tau) expected_increment = perturbation / 12 + 11 * perturbation ** 2 / 60 assert abs(perturbed_cost - base_cost - expected_increment) < 1e-9 # A feasible clamped-endpoint perturbation verifies the global energy gap. for distance, duration, amplitude in ((1, 2, 1), (-2, 3, -0.5), (0, 1, 2)): t = duration * tau _, _, optimal_u = minimum_energy_transfer(distance, duration, t) delta_u = amplitude / duration ** 2 * (2 - 12 * tau + 12 * tau ** 2) energy_gap = np.trapezoid((optimal_u + delta_u) ** 2 - optimal_u ** 2, t) exact_gap = 4 * amplitude ** 2 / (5 * duration ** 3) assert abs(energy_gap - exact_gap) < 1e-7 # Active-set regimes and the sign of the upper-bound multiplier. for bound, expected_mu in ((0.25, 2), (0.5, 0), (1, 0)): acceleration = two_pulse_transfer(1, 2, penalty=2, limit=bound)[0] upper_mu = -(8 * acceleration - 4) assert abs(upper_mu - expected_mu) < 1e-12 assert upper_mu >= 0 and abs(upper_mu * (acceleration - bound)) < 1e-12 # Lecture 4: the terminal Hamiltonian condition agrees with scalar minimization. optimal_duration, free_time_cost = free_time_transfer(-10) assert abs(optimal_duration - 1800 ** 0.2) < 1e-12 d, alpha, beta = -10., 1., 1. t = np.linspace(0, optimal_duration, 10001) _, velocity, acceleration = minimum_energy_transfer(d, optimal_duration, t) position_costate = -12 * beta * d / optimal_duration ** 3 velocity_costate = -beta * acceleration hamiltonian = (beta * acceleration ** 2 / 2 + position_costate * velocity + velocity_costate * acceleration) assert np.ptp(hamiltonian) < 1e-12 assert abs(hamiltonian[-1] + alpha * optimal_duration) < 1e-12 for scale in (0.8, 0.99, 1.01, 1.2): trial_t = scale * optimal_duration trial_cost = alpha * trial_t ** 2 / 2 + 6 * beta * d ** 2 / trial_t ** 3 assert trial_cost > free_time_cost assert np.allclose(free_time_transfer(10), free_time_transfer(-10)) for invalid in ((0, 1, 1), (1, 0, 1), (1, 1, 0)): try: free_time_transfer(*invalid) except ValueError: pass else: raise AssertionError('Degenerate free-time problem was accepted') # Lecture 5: endpoint moments check the re-solved bounded profile. bounded_time = np.linspace(0, 2.2, 100001) bounded_u, bounded_effort, ramp_width = bounded_energy_profile(1, 2.2, 1, bounded_time) assert np.max(abs(bounded_u)) <= 1 assert abs(np.trapezoid(bounded_u, bounded_time)) < 1e-12 assert abs(np.trapezoid((2.2 - bounded_time) * bounded_u, bounded_time) - 1) < 1e-8 assert abs(np.trapezoid(bounded_u ** 2, bounded_time) - bounded_effort) < 1e-8 negative_u, negative_effort, _ = bounded_energy_profile(-1, 2.2, 1, bounded_time) assert np.allclose(negative_u, -bounded_u) and negative_effort == bounded_effort assert np.allclose(bounded_energy_profile(0, 2.2, 1, bounded_time)[0], 0) assert abs(bounded_energy_profile(1, 2, 1, np.array([0., 1., 2.]))[1] - 2) < 1e-12 try: bounded_energy_profile(1, 1.9, 1, np.array([0., 1.9])) except ValueError: pass else: raise AssertionError('Infeasible bounded energy transfer accepted') trial_inputs = np.linspace(-1, 1, 20001) for sigma in (-3., -1., 0., 1., 3.): exact_energy_input = np.clip(-sigma, -1, 1) assert (exact_energy_input ** 2 / 2 + sigma * exact_energy_input <= np.min(trial_inputs ** 2 / 2 + sigma * trial_inputs) + 1e-12) exact_fuel_input = -1 if sigma > 1 else 1 if sigma < -1 else 0 assert (abs(exact_fuel_input) + sigma * exact_fuel_input <= np.min(abs(trial_inputs) + sigma * trial_inputs) + 1e-12) clipped = np.clip(6 / 2.2 ** 2 * (1 - 2 * bounded_time / 2.2), -1, 1) clipped_position = np.trapezoid((2.2 - bounded_time) * clipped, bounded_time) assert abs(clipped_position - (1.1 ** 2 - (2.2 ** 3 / 12) ** 2 / 3)) < 1e-8 from scipy.integrate import solve_bvp def bvp_rhs(t, z): return np.vstack((z[1], -np.abs(z[0]))) def bvp_boundary(za, zb): return np.array([za[0], zb[0] + 2]) bvp_mesh = np.linspace(0, 4, 41) dense = np.linspace(0, 4, 1001) bvp_errors = {} for branch in ('negative', 'positive'): if branch == 'negative': guess = np.vstack((-bvp_mesh / 2, -np.ones_like(bvp_mesh) / 2)) exact = -2 * np.sinh(dense) / np.sinh(4) else: guess = np.vstack((2 * np.sin(bvp_mesh), 2 * np.cos(bvp_mesh))) amplitude = 2 / np.sinh(4 - np.pi) exact = np.where(dense <= np.pi, amplitude * np.sin(dense), -amplitude * np.sinh(dense - np.pi)) result = solve_bvp(bvp_rhs, bvp_boundary, bvp_mesh, guess, tol=1e-7, max_nodes=5000) assert result.success, result.message evaluated = result.sol(dense) error = float(np.max(abs(evaluated[0] - exact))) assert error < 1e-6 assert np.max(abs(bvp_boundary(evaluated[:, 0], evaluated[:, -1]))) < 1e-8 assert np.max(np.diff(evaluated[1])) < 1e-7 bvp_errors[branch] = error # Riccati recursion is checked against an independently condensed QP. a = np.array([[1., 0.2], [0., 1.]]) b = np.array([[0.02], [0.2]]) q, r, terminal = np.diag([1., 0.1]), np.array([[0.05]]), np.eye(2) initial, horizon = np.array([1., 0.]), 12 gains, values = finite_lqr(a, b, q, r, terminal, horizon) states, actions = rollout(a, b, gains, initial) phi = np.vstack([np.linalg.matrix_power(a, k) for k in range(horizon + 1)]) gamma = np.zeros((2 * (horizon + 1), horizon)) for k in range(1, horizon + 1): for j in range(k): gamma[2 * k:2 * k + 2, j:j + 1] = ( np.linalg.matrix_power(a, k - 1 - j) @ b) weight = np.kron(np.eye(horizon + 1), q) weight[-2:, -2:] = terminal dense_h = gamma.T @ weight @ gamma + r[0, 0] * np.eye(horizon) dense_u = -np.linalg.solve(dense_h, gamma.T @ weight @ phi @ initial) assert np.allclose(actions[:, 0], dense_u, atol=1e-11) actual_cost = sum(z @ q @ z for z in states[:-1]) actual_cost += float(np.sum(actions ** 2) * r[0, 0]) actual_cost += states[-1] @ terminal @ states[-1] assert abs(actual_cost - initial @ values[0] @ initial) < 1e-11 zero_gains, zero_values = finite_lqr(a, b, q, r, terminal, 0) assert len(zero_gains) == 0 and np.allclose(zero_values[0], terminal) # Nonlinear local optimization: actual rollout, exact discrete adjoint, # and the zero-drag reduction to the independently implemented LQR. nonlinear_states, nonlinear_actions, accepted_costs, nonlinear_gradient = ( nonlinear_cart_ilqr()) assert all(v < u for u, v in zip(accepted_costs, accepted_costs[1:])) assert nonlinear_gradient < 1e-7 linear_states, linear_actions, _, _ = nonlinear_cart_ilqr(drag=0) linear_a, linear_b = np.array([[1., .1], [0., 1.]]), np.array([[0.], [.1]]) linear_gains, _ = finite_lqr( linear_a, linear_b, np.diag([1., .1]), np.array([[.05]]), np.diag([20., 2.]), 30) expected_states, expected_actions = rollout( linear_a, linear_b, linear_gains, np.array([1., 0.])) assert np.max(abs(linear_actions - expected_actions[:, 0])) < 1e-9 assert np.max(abs(linear_states - expected_states)) < 1e-9 for k in range(30): position, velocity = nonlinear_states[k] expected_next = np.array([ position + .1*velocity, velocity + .1*(nonlinear_actions[k] - .2*velocity**3)]) assert np.max(abs(expected_next - nonlinear_states[k+1])) < 1e-13 for invalid in ({'horizon': 0}, {'drag': -1}, {'initial': [1]}): try: nonlinear_cart_ilqr(**invalid) except ValueError: pass else: raise AssertionError('Invalid iLQR problem accepted') # Full DDP includes the dynamics curvature contracted with value gradient. def scalar_nonlinear_cost(action): successor = action + .5*action**2 return .05*action**2 + .5*(successor - 1)**2 delta = 1e-4 curvature = (scalar_nonlinear_cost(delta) - 2*scalar_nonlinear_cost(0) + scalar_nonlinear_cost(-delta))/delta**2 assert abs(curvature - .1) < 1e-7 assert scalar_nonlinear_cost(10) > scalar_nonlinear_cost(0) assert scalar_nonlinear_cost(10/16) < scalar_nonlinear_cost(0) # A two-state model with an exact policy-evaluation solution. transition = np.array([[[1, 0], [0, 1]], [[0, 1], [0, 1]]]) costs = np.array([[1, 0.2], [0, 0]]) value, policy, bound = discounted_value_iteration(transition, costs, 0.9) assert np.allclose(value, [0.2, 0]) and policy[0] == 1 assert np.max(abs(value - [0.2, 0])) <= bound + 1e-14 immediate, _, _ = discounted_value_iteration(transition, costs, 0) assert np.allclose(immediate, costs.min(axis=1)) pred = robust_predecessor( range(-3, 4), [-1, 0, 1], [-1, 0, 1], {-1, 0, 1}, lambda x, u, w: x + u + w) assert pred == {-1, 0, 1} no_disturbance = robust_predecessor( range(-3, 4), [-1, 0, 1], [0], {-1, 0, 1}, lambda x, u, w: x + u + w) assert no_disturbance == {-2, -1, 0, 1, 2} rng = np.random.default_rng(203) samples, controls = rng.normal(size=(40, 2)), rng.normal(size=(40, 1)) next_samples = samples @ a.T + controls @ b.T fit_a, fit_b, residual = identify_linear(samples, controls, next_samples) assert np.allclose(fit_a, a) and np.allclose(fit_b, b) try: identify_linear(samples, samples[:, :1], next_samples) except ValueError: pass else: raise AssertionError('Rank-deficient identification was accepted') mean, gradient = bernoulli_policy_gradient(0.4, -1, 2) eps = 1e-5 plus = bernoulli_policy_gradient(0.4 + eps, -1, 2)[0] minus = bernoulli_policy_gradient(0.4 - eps, -1, 2)[0] assert abs((plus - minus) / (2 * eps) - gradient) < 1e-9 assert bernoulli_policy_gradient(1000, -1, 2) == (2., 0.) assert bernoulli_policy_gradient(-1000, -1, 2) == (-1., 0.) return {'two_pulse_hard': list(hard), 'two_pulse_soft': list(soft), 'pulse_equality_inputs': equality_actions.tolist(), 'pulse_equality_multiplier': equality_multiplier, 'pulse_value_sensitivity': float(sensitivity), 'armijo_step': list(step_result), 'first_variation': float(first_variation), 'variation_increment_and_energy_gap_checks': 'passed', 'active_set_multiplier_checks': 'passed', 'free_time_duration': optimal_duration, 'free_time_cost': free_time_cost, 'terminal_hamiltonian_residual': float( hamiltonian[-1] + alpha * optimal_duration), 'bounded_energy_effort': bounded_effort, 'bounded_energy_ramp_half_width': ramp_width, 'incorrectly_clipped_position': float(clipped_position), 'bvp_branch_position_errors': bvp_errors, 'minimum_energy': float(energy), 'lqr_cost': float(actual_cost), 'lqr_dense_max_error': float( np.max(abs(actions[:, 0] - dense_u))), 'ilqr_accepted_costs': accepted_costs, 'ilqr_gradient_infinity_norm': nonlinear_gradient, 'ilqr_terminal_state': nonlinear_states[-1].tolist(), 'ilqr_linear_limit_max_action_error': float( np.max(abs(linear_actions - expected_actions[:, 0]))), 'ddp_scalar_curvature_finite_difference': float(curvature), 'bellman_value': value.tolist(), 'bellman_error_bound': bound, 'robust_predecessor': sorted(pred), 'identification_residual': residual, 'policy_reward': mean, 'policy_gradient': gradient, 'checks': 'passed'} if __name__ == '__main__': print(json.dumps(run_checks(), indent=2))