Add files using upload-large-folder tool

f8faa82 verified 18 days ago

17.5 kB

"""
IpoptProjector: projects a reference trajectory onto a feasible set
defined by circular/planar obstacles and optional linear dynamics.
CasADi uses column-major (Fortran) order for matrices.
All constraint geometry is in real-world (unnormalized) coordinates;
the decision variable X is in normalized space.
"""

from __future__ import annotations

import copy
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import Callable, Dict, List, Optional, Tuple

import casadi as cs
import numpy as np

from obstacles import CircularObstacle, PlanarObstacle


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _build_and_solve(
    config: dict,
    X_ref: np.ndarray,
    s0: np.ndarray,
    reg_weight: float = None,
    a_prev: np.ndarray = None,
) -> Tuple[np.ndarray, str, int]:
    """Standalone function for parallel workers (must be picklable)."""
    proj = IpoptProjector(**config)
    return proj.project(X_ref, s0, reg_weight=reg_weight, a_prev=a_prev)


# ---------------------------------------------------------------------------
# IpoptProjector
# ---------------------------------------------------------------------------

class IpoptProjector:

    def __init__(
        self,
        circular_obstacles: List[CircularObstacle],
        planar_obstacles: List[PlanarObstacle],
        horizon: int,
        action_dim: int,
        state_dim: int,
        normalizer,
        dynamics: Optional[Dict[str, np.ndarray]] = None,
        dynamics_relaxation: float = 0.0,
        obstacle_margin: float = 0.0,
        input_saturation: bool = False,
        objective_fn: Optional[Callable] = None,
        objective_scale: float = 1.0,
        max_action_delta: float = 0.0,
        solver_opts: Optional[dict] = None,
    ):
        self.circular_obstacles = circular_obstacles
        self.planar_obstacles = planar_obstacles
        self.horizon = horizon
        self.action_dim = action_dim
        self.state_dim = state_dim
        self.transition_dim = action_dim + state_dim
        self.normalizer = normalizer
        self.dynamics = dynamics
        self.dynamics_relaxation = dynamics_relaxation
        self.obstacle_margin = obstacle_margin
        self.input_saturation = input_saturation
        self.objective_fn = objective_fn
        self.objective_scale = objective_scale
        self.max_action_delta = max_action_delta

        # DOF: full trajectory minus the fixed initial state
        # layout: [a_0, | a_1, s_1, | a_2, s_2, | ... | a_{H-1}, s_{H-1}]
        self.dof = horizon * self.transition_dim - state_dim

        # normalizer bounds for x_pos (dim 2) and y_pos (dim 3) of observations
        obs_mins = normalizer.normalizers["observations"].mins
        obs_maxs = normalizer.normalizers["observations"].maxs
        self._xmin = float(obs_mins[2])
        self._xmax = float(obs_maxs[2])
        self._ymin = float(obs_mins[3])
        self._ymax = float(obs_maxs[3])

        # action bounds in normalized space ([-1, 1] for LimitsNormalizer)
        act_mins = normalizer.normalizers["actions"].mins
        act_maxs = normalizer.normalizers["actions"].maxs
        self._act_norm_lo = np.array([
            (lo - lo) / (hi - lo) * 2 - 1 for lo, hi in zip(act_mins, act_maxs)
        ])  # always -1
        self._act_norm_hi = np.array([
            (hi - lo) / (hi - lo) * 2 - 1 for lo, hi in zip(act_mins, act_maxs)
        ])  # always +1

        # precompute index arrays into the flat DOF vector
        self._precompute_indices()

        # solver options (IPOPT defaults, output silencing only)
        self._solver_opts = {
            "ipopt.print_level": 0,
            "ipopt.sb": "yes",
            "print_time": 0,
        }
        if solver_opts is not None:
            self._solver_opts.update(solver_opts)

        # build the CasADi problem
        self._formulate()

    # ------------------------------------------------------------------
    # Index precomputation
    # ------------------------------------------------------------------

    def _precompute_indices(self):
        """Compute indices into the flat DOF vector for states and actions.
        DOF layout:
          [a_0(action_dim) | a_1(action_dim), s_1(state_dim) | ... | a_{H-1}(action_dim), s_{H-1}(state_dim)]
        s_0 is NOT in the DOF — it is a parameter.
        """
        ad = self.action_dim
        sd = self.state_dim
        td = self.transition_dim

        # state indices for s_1 .. s_{H-1}  (H-1 states in DOF)
        # first state s_1 starts at offset ad (for a_0) + ad (for a_1) = 2*ad
        self._state_indices = []      # list of (start, end) for each state block
        self._xpos_indices = []       # index of x_pos (obs dim 2) for each state
        self._ypos_indices = []       # index of y_pos (obs dim 3) for each state

        for t in range(1, self.horizon):
            # block t: starts at ad + (t-1)*td + ad = ad + (t-1)*td + ad
            # actually: block 0 is just a_0 (ad elements)
            #           block t (t>=1) is a_t (ad) + s_t (sd), starting at ad + (t-1)*td
            block_start = ad + (t - 1) * td
            s_start = block_start + ad  # state starts after action
            s_end = s_start + sd

            self._state_indices.append((s_start, s_end))
            self._xpos_indices.append(s_start + 2)  # obs dim 2 = x_pos
            self._ypos_indices.append(s_start + 3)  # obs dim 3 = y_pos

        # action indices for a_0 .. a_{H-1}
        self._action_indices = []
        # a_0 at offset 0
        self._action_indices.append((0, ad))
        for t in range(1, self.horizon):
            block_start = ad + (t - 1) * td
            self._action_indices.append((block_start, block_start + ad))

    # ------------------------------------------------------------------
    # Unnormalization (normalized [-1,1] -> real-world)
    # ------------------------------------------------------------------

    def _unnorm_x(self, x_norm):
        return (x_norm + 1.0) / 2.0 * (self._xmax - self._xmin) + self._xmin

    def _unnorm_y(self, y_norm):
        return (y_norm + 1.0) / 2.0 * (self._ymax - self._ymin) + self._ymin

    # ------------------------------------------------------------------
    # Problem formulation
    # ------------------------------------------------------------------

    def _formulate(self):
        opti = cs.Opti()

        # decision variable (column vector, CasADi column-major)
        X = opti.variable(self.dof, 1)
        X_ref = opti.parameter(self.dof, 1)
        s0 = opti.parameter(self.state_dim, 1)
        reg_weight = opti.parameter(1, 1)
        opti.set_value(reg_weight, 1.0)

        # --- objective: projection ---
        cost = 0.5 * reg_weight * cs.sumsqr(X - X_ref)

        if self.objective_fn is not None:
            cost = cost + self.objective_scale * self.objective_fn(X, s0)

        opti.minimize(cost)

        # --- circular obstacle constraints (squared distance) ---
        for obs in self.circular_obstacles:
            cx, cy = obs.center
            r_eff_sq = (obs.radius + self.obstacle_margin) ** 2

            for xi, yi in zip(self._xpos_indices, self._ypos_indices):
                xu = self._unnorm_x(X[xi])
                yu = self._unnorm_y(X[yi])
                opti.subject_to((xu - cx) ** 2 + (yu - cy) ** 2 >= r_eff_sq)

        # --- planar obstacle constraints ---
        for pobs in self.planar_obstacles:
            slope = pobs.slope
            intercept = pobs.intercept

            for xi, yi in zip(self._xpos_indices, self._ypos_indices):
                xu = self._unnorm_x(X[xi])
                yu = self._unnorm_y(X[yi])
                opti.subject_to(slope * xu + intercept - yu >= self.obstacle_margin)

        # --- input saturation (action box constraints in normalized space) ---
        if self.input_saturation:
            for a_start, a_end in self._action_indices:
                for d in range(self.action_dim):
                    opti.subject_to(X[a_start + d] >= self._act_norm_lo[d])
                    opti.subject_to(X[a_start + d] <= self._act_norm_hi[d])

        # --- dynamics constraints (optional) ---
        if self.dynamics is not None:
            A = self.dynamics["A"]  # (state_dim, state_dim)
            B = self.dynamics["B"]  # (state_dim, action_dim)
            c_vec = self.dynamics["c"]  # (state_dim,)

            # CasADi is column-major: use DM for constant matrices
            A_cs = cs.DM(A)    # (state_dim x state_dim)
            B_cs = cs.DM(B)    # (state_dim x action_dim)
            c_cs = cs.DM(c_vec.reshape(-1, 1))  # (state_dim x 1)

            relax = self.dynamics_relaxation

            # t=0: s_1 = A @ s_0 + B @ a_0 + c  (s_0 is parameter)
            a0_start, a0_end = self._action_indices[0]
            s1_start, s1_end = self._state_indices[0]
            a_0 = X[a0_start:a0_end]
            s_1 = X[s1_start:s1_end]
            pred_0 = cs.mtimes(A_cs, s0) + cs.mtimes(B_cs, a_0) + c_cs
            if relax > 0:
                opti.subject_to(s_1 <= pred_0 + relax)
                opti.subject_to(s_1 >= pred_0 - relax)
            else:
                opti.subject_to(s_1 == pred_0)

            # t=1..H-2: s_{t+1} = A @ s_t + B @ a_t + c
            for t in range(1, self.horizon - 1):
                s_t_start, s_t_end = self._state_indices[t - 1]
                a_t_start, a_t_end = self._action_indices[t]

                if t >= len(self._state_indices):
                    break

                s_t = X[s_t_start:s_t_end]
                a_t = X[a_t_start:a_t_end]
                sn_start, sn_end = self._state_indices[t]
                s_next = X[sn_start:sn_end]
                pred_t = cs.mtimes(A_cs, s_t) + cs.mtimes(B_cs, a_t) + c_cs
                if relax > 0:
                    opti.subject_to(s_next <= pred_t + relax)
                    opti.subject_to(s_next >= pred_t - relax)
                else:
                    opti.subject_to(s_next == pred_t)

        self._a_prev = None

        # --- solver ---
        opti.solver("ipopt", self._solver_opts)

        # store references
        self._opti = opti
        self._X = X
        self._X_ref = X_ref
        self._s0 = s0
        self._reg_weight = reg_weight

    # ------------------------------------------------------------------
    # Solve
    # ------------------------------------------------------------------

    def set_regularization_weight(self, w: float):
        """Set the regularization weight for the tracking cost."""
        self._opti.set_value(self._reg_weight, w)

    def project(self, X_ref: np.ndarray, s0: np.ndarray, reg_weight: float = None, a_prev: np.ndarray = None) -> Tuple[np.ndarray, str, int]:
        """Project a single reference trajectory onto the feasible set.
        Args:
            X_ref: (dof,) flat DOF vector in normalized space.
            s0: (state_dim,) initial state in normalized space.
            reg_weight: optional regularization weight (overrides default 1.0).
            a_prev: (action_dim,) last executed action in normalized space (for action rate constraint).
        Returns:
            X_proj: (dof,) projected DOF vector.
            status: "optimal" if the solver converged, "failed" otherwise.
            n_iters: number of IPOPT iterations spent on the solve.
        """
        X_ref = np.asarray(X_ref).flatten()
        s0 = np.asarray(s0).flatten()

        assert X_ref.shape == (self.dof,), f"X_ref shape {X_ref.shape} != ({self.dof},)"
        assert s0.shape == (self.state_dim,), f"s0 shape {s0.shape} != ({self.state_dim},)"

        # CasADi parameters are column vectors
        self._opti.set_value(self._X_ref, X_ref.reshape(-1, 1))
        self._opti.set_value(self._s0, s0.reshape(-1, 1))
        if reg_weight is not None:
            self._opti.set_value(self._reg_weight, reg_weight)
        if a_prev is not None and self._a_prev is not None:
            self._opti.set_value(self._a_prev, np.asarray(a_prev).reshape(-1, 1))
        self._opti.set_initial(self._X, X_ref.reshape(-1, 1))

        try:
            sol = self._opti.solve()
            X_proj = np.array(sol.value(self._X)).flatten()
            status = "optimal"
            stats_src = self._opti
        except RuntimeError:
            X_proj = np.array(self._opti.debug.value(self._X)).flatten()
            status = "failed"
            stats_src = self._opti.debug

        try:
            n_iters = int(stats_src.stats()["iter_count"])
        except Exception:
            n_iters = 0

        return X_proj, status, n_iters

    def project_batch(
        self,
        X_refs: np.ndarray,
        s0s: np.ndarray,
        n_workers: int = 1,
        reg_weight: float = None,
        a_prevs: np.ndarray = None,
    ) -> Tuple[np.ndarray, List[str], List[int]]:
        """Project a batch of trajectories in parallel.
        Args:
            X_refs: (batch, dof)
            s0s: (batch, state_dim)
            n_workers: number of parallel IPOPT solves.
            reg_weight: optional regularization weight (passed to each solve).
            a_prevs: (batch, action_dim) last executed actions for action rate constraint.
        Returns:
            X_projs: (batch, dof)
            statuses: list of per-solve statuses ("optimal"/"failed").
            iters: list of per-solve IPOPT iteration counts.
        """
        batch_size = X_refs.shape[0]
        assert X_refs.shape == (batch_size, self.dof)
        assert s0s.shape == (batch_size, self.state_dim)

        if n_workers <= 1:
            projs = []
            statuses = []
            iters = []
            for i in range(batch_size):
                ap = a_prevs[i] if a_prevs is not None else None
                x_proj, status, n_iters = self.project(X_refs[i], s0s[i], reg_weight=reg_weight, a_prev=ap)
                projs.append(x_proj)
                statuses.append(status)
                iters.append(n_iters)
            return np.stack(projs), statuses, iters

        # parallel: build config dict for picklable worker function
        config = {
            "circular_obstacles": self.circular_obstacles,
            "planar_obstacles": self.planar_obstacles,
            "horizon": self.horizon,
            "action_dim": self.action_dim,
            "state_dim": self.state_dim,
            "normalizer": self.normalizer,
            "dynamics": self.dynamics,
            "dynamics_relaxation": self.dynamics_relaxation,
            "obstacle_margin": self.obstacle_margin,
            "input_saturation": self.input_saturation,
            "objective_fn": self.objective_fn,
            "objective_scale": self.objective_scale,
            "max_action_delta": self.max_action_delta,
            "solver_opts": self._solver_opts,
        }

        results = [None] * batch_size
        statuses = [None] * batch_size
        iters = [0] * batch_size
        with ProcessPoolExecutor(max_workers=n_workers) as pool:
            futures = {
                pool.submit(
                    _build_and_solve, config, X_refs[i], s0s[i], reg_weight,
                    a_prevs[i] if a_prevs is not None else None,
                ): i
                for i in range(batch_size)
            }
            for future in as_completed(futures):
                idx = futures[future]
                x_proj, status, n_iters = future.result()
                results[idx] = x_proj
                statuses[idx] = status
                iters[idx] = n_iters

        return np.stack(results), statuses, iters

    # ------------------------------------------------------------------
    # Utilities
    # ------------------------------------------------------------------

    def dof_to_full_trajectory(self, X_dof: np.ndarray, s0: np.ndarray) -> np.ndarray:
        """Reconstruct the full (horizon, transition_dim) trajectory from DOF + s0.
        Returns:
            traj: (horizon, action_dim + state_dim)
        """
        X_dof = np.asarray(X_dof).flatten()
        s0 = np.asarray(s0).flatten()
        ad = self.action_dim
        sd = self.state_dim

        traj = np.zeros((self.horizon, ad + sd))

        # timestep 0: action from DOF, state = s0
        a0_s, a0_e = self._action_indices[0]
        traj[0, :ad] = X_dof[a0_s:a0_e]
        traj[0, ad:] = s0

        # timesteps 1..H-1: action + state from DOF
        for t in range(1, self.horizon):
            a_s, a_e = self._action_indices[t]
            s_s, s_e = self._state_indices[t - 1]
            traj[t, :ad] = X_dof[a_s:a_e]
            traj[t, ad:] = X_dof[s_s:s_e]

        return traj

    def full_trajectory_to_dof(self, traj: np.ndarray) -> np.ndarray:
        """Extract the DOF vector from a full (horizon, transition_dim) trajectory.
        s_0 is dropped (it's a parameter, not a DOF).
        """
        ad = self.action_dim
        sd = self.state_dim
        X_dof = np.zeros(self.dof)

        # a_0
        a0_s, a0_e = self._action_indices[0]
        X_dof[a0_s:a0_e] = traj[0, :ad]

        # a_t, s_t for t >= 1
        for t in range(1, self.horizon):
            a_s, a_e = self._action_indices[t]
            s_s, s_e = self._state_indices[t - 1]
            X_dof[a_s:a_e] = traj[t, :ad]
            X_dof[s_s:s_e] = traj[t, ad:]

        return X_dof

	"""
	IpoptProjector: projects a reference trajectory onto a feasible set
	defined by circular/planar obstacles and optional linear dynamics.

	CasADi uses column-major (Fortran) order for matrices.
	All constraint geometry is in real-world (unnormalized) coordinates;
	the decision variable X is in normalized space.
	"""

	from __future__ import annotations

	import copy
	from concurrent.futures import ProcessPoolExecutor, as_completed
	from typing import Callable, Dict, List, Optional, Tuple

	import casadi as cs
	import numpy as np

	from obstacles import CircularObstacle, PlanarObstacle


	# ---------------------------------------------------------------------------
	# Helpers
	# ---------------------------------------------------------------------------

	def _build_and_solve(
	config: dict,
	X_ref: np.ndarray,
	s0: np.ndarray,
	reg_weight: float = None,
	a_prev: np.ndarray = None,
	) -> Tuple[np.ndarray, str, int]:
	"""Standalone function for parallel workers (must be picklable)."""
	proj = IpoptProjector(**config)
	return proj.project(X_ref, s0, reg_weight=reg_weight, a_prev=a_prev)


	# ---------------------------------------------------------------------------
	# IpoptProjector
	# ---------------------------------------------------------------------------

	class IpoptProjector:

	def __init__(
	self,
	circular_obstacles: List[CircularObstacle],
	planar_obstacles: List[PlanarObstacle],
	horizon: int,
	action_dim: int,
	state_dim: int,
	normalizer,
	dynamics: Optional[Dict[str, np.ndarray]] = None,
	dynamics_relaxation: float = 0.0,
	obstacle_margin: float = 0.0,
	input_saturation: bool = False,
	objective_fn: Optional[Callable] = None,
	objective_scale: float = 1.0,
	max_action_delta: float = 0.0,
	solver_opts: Optional[dict] = None,
	):
	self.circular_obstacles = circular_obstacles
	self.planar_obstacles = planar_obstacles
	self.horizon = horizon
	self.action_dim = action_dim
	self.state_dim = state_dim
	self.transition_dim = action_dim + state_dim
	self.normalizer = normalizer
	self.dynamics = dynamics
	self.dynamics_relaxation = dynamics_relaxation
	self.obstacle_margin = obstacle_margin
	self.input_saturation = input_saturation
	self.objective_fn = objective_fn
	self.objective_scale = objective_scale
	self.max_action_delta = max_action_delta

	# DOF: full trajectory minus the fixed initial state
	# layout: [a_0, \| a_1, s_1, \| a_2, s_2, \| ... \| a_{H-1}, s_{H-1}]
	self.dof = horizon * self.transition_dim - state_dim

	# normalizer bounds for x_pos (dim 2) and y_pos (dim 3) of observations
	obs_mins = normalizer.normalizers["observations"].mins
	obs_maxs = normalizer.normalizers["observations"].maxs
	self._xmin = float(obs_mins[2])
	self._xmax = float(obs_maxs[2])
	self._ymin = float(obs_mins[3])
	self._ymax = float(obs_maxs[3])

	# action bounds in normalized space ([-1, 1] for LimitsNormalizer)
	act_mins = normalizer.normalizers["actions"].mins
	act_maxs = normalizer.normalizers["actions"].maxs
	self._act_norm_lo = np.array([
	(lo - lo) / (hi - lo) * 2 - 1 for lo, hi in zip(act_mins, act_maxs)
	]) # always -1
	self._act_norm_hi = np.array([
	(hi - lo) / (hi - lo) * 2 - 1 for lo, hi in zip(act_mins, act_maxs)
	]) # always +1

	# precompute index arrays into the flat DOF vector
	self._precompute_indices()

	# solver options (IPOPT defaults, output silencing only)
	self._solver_opts = {
	"ipopt.print_level": 0,
	"ipopt.sb": "yes",
	"print_time": 0,
	}
	if solver_opts is not None:
	self._solver_opts.update(solver_opts)

	# build the CasADi problem
	self._formulate()

	# ------------------------------------------------------------------
	# Index precomputation
	# ------------------------------------------------------------------

	def _precompute_indices(self):
	"""Compute indices into the flat DOF vector for states and actions.

	DOF layout:
	[a_0(action_dim) \| a_1(action_dim), s_1(state_dim) \| ... \| a_{H-1}(action_dim), s_{H-1}(state_dim)]

	s_0 is NOT in the DOF — it is a parameter.
	"""
	ad = self.action_dim
	sd = self.state_dim
	td = self.transition_dim

	# state indices for s_1 .. s_{H-1} (H-1 states in DOF)
	# first state s_1 starts at offset ad (for a_0) + ad (for a_1) = 2*ad
	self._state_indices = [] # list of (start, end) for each state block
	self._xpos_indices = [] # index of x_pos (obs dim 2) for each state
	self._ypos_indices = [] # index of y_pos (obs dim 3) for each state

	for t in range(1, self.horizon):
	# block t: starts at ad + (t-1)td + ad = ad + (t-1)td + ad
	# actually: block 0 is just a_0 (ad elements)
	# block t (t>=1) is a_t (ad) + s_t (sd), starting at ad + (t-1)*td
	block_start = ad + (t - 1) * td
	s_start = block_start + ad # state starts after action
	s_end = s_start + sd

	self._state_indices.append((s_start, s_end))
	self._xpos_indices.append(s_start + 2) # obs dim 2 = x_pos
	self._ypos_indices.append(s_start + 3) # obs dim 3 = y_pos

	# action indices for a_0 .. a_{H-1}
	self._action_indices = []
	# a_0 at offset 0
	self._action_indices.append((0, ad))
	for t in range(1, self.horizon):
	block_start = ad + (t - 1) * td
	self._action_indices.append((block_start, block_start + ad))

	# ------------------------------------------------------------------
	# Unnormalization (normalized [-1,1] -> real-world)
	# ------------------------------------------------------------------

	def _unnorm_x(self, x_norm):
	return (x_norm + 1.0) / 2.0 * (self._xmax - self._xmin) + self._xmin

	def _unnorm_y(self, y_norm):
	return (y_norm + 1.0) / 2.0 * (self._ymax - self._ymin) + self._ymin

	# ------------------------------------------------------------------
	# Problem formulation
	# ------------------------------------------------------------------

	def _formulate(self):
	opti = cs.Opti()

	# decision variable (column vector, CasADi column-major)
	X = opti.variable(self.dof, 1)
	X_ref = opti.parameter(self.dof, 1)
	s0 = opti.parameter(self.state_dim, 1)
	reg_weight = opti.parameter(1, 1)
	opti.set_value(reg_weight, 1.0)

	# --- objective: projection ---
	cost = 0.5 * reg_weight * cs.sumsqr(X - X_ref)

	if self.objective_fn is not None:
	cost = cost + self.objective_scale * self.objective_fn(X, s0)

	opti.minimize(cost)

	# --- circular obstacle constraints (squared distance) ---
	for obs in self.circular_obstacles:
	cx, cy = obs.center
	r_eff_sq = (obs.radius + self.obstacle_margin) ** 2

	for xi, yi in zip(self._xpos_indices, self._ypos_indices):
	xu = self._unnorm_x(X[xi])
	yu = self._unnorm_y(X[yi])
	opti.subject_to((xu - cx) 2 + (yu - cy) 2 >= r_eff_sq)

	# --- planar obstacle constraints ---
	for pobs in self.planar_obstacles:
	slope = pobs.slope
	intercept = pobs.intercept

	for xi, yi in zip(self._xpos_indices, self._ypos_indices):
	xu = self._unnorm_x(X[xi])
	yu = self._unnorm_y(X[yi])
	opti.subject_to(slope * xu + intercept - yu >= self.obstacle_margin)

	# --- input saturation (action box constraints in normalized space) ---
	if self.input_saturation:
	for a_start, a_end in self._action_indices:
	for d in range(self.action_dim):
	opti.subject_to(X[a_start + d] >= self._act_norm_lo[d])
	opti.subject_to(X[a_start + d] <= self._act_norm_hi[d])

	# --- dynamics constraints (optional) ---
	if self.dynamics is not None:
	A = self.dynamics["A"] # (state_dim, state_dim)
	B = self.dynamics["B"] # (state_dim, action_dim)
	c_vec = self.dynamics["c"] # (state_dim,)

	# CasADi is column-major: use DM for constant matrices
	A_cs = cs.DM(A) # (state_dim x state_dim)
	B_cs = cs.DM(B) # (state_dim x action_dim)
	c_cs = cs.DM(c_vec.reshape(-1, 1)) # (state_dim x 1)

	relax = self.dynamics_relaxation

	# t=0: s_1 = A @ s_0 + B @ a_0 + c (s_0 is parameter)
	a0_start, a0_end = self._action_indices[0]
	s1_start, s1_end = self._state_indices[0]
	a_0 = X[a0_start:a0_end]
	s_1 = X[s1_start:s1_end]
	pred_0 = cs.mtimes(A_cs, s0) + cs.mtimes(B_cs, a_0) + c_cs
	if relax > 0:
	opti.subject_to(s_1 <= pred_0 + relax)
	opti.subject_to(s_1 >= pred_0 - relax)
	else:
	opti.subject_to(s_1 == pred_0)

	# t=1..H-2: s_{t+1} = A @ s_t + B @ a_t + c
	for t in range(1, self.horizon - 1):
	s_t_start, s_t_end = self._state_indices[t - 1]
	a_t_start, a_t_end = self._action_indices[t]

	if t >= len(self._state_indices):
	break

	s_t = X[s_t_start:s_t_end]
	a_t = X[a_t_start:a_t_end]
	sn_start, sn_end = self._state_indices[t]
	s_next = X[sn_start:sn_end]
	pred_t = cs.mtimes(A_cs, s_t) + cs.mtimes(B_cs, a_t) + c_cs
	if relax > 0:
	opti.subject_to(s_next <= pred_t + relax)
	opti.subject_to(s_next >= pred_t - relax)
	else:
	opti.subject_to(s_next == pred_t)

	self._a_prev = None

	# --- solver ---
	opti.solver("ipopt", self._solver_opts)

	# store references
	self._opti = opti
	self._X = X
	self._X_ref = X_ref
	self._s0 = s0
	self._reg_weight = reg_weight

	# ------------------------------------------------------------------
	# Solve
	# ------------------------------------------------------------------

	def set_regularization_weight(self, w: float):
	"""Set the regularization weight for the tracking cost."""
	self._opti.set_value(self._reg_weight, w)

	def project(self, X_ref: np.ndarray, s0: np.ndarray, reg_weight: float = None, a_prev: np.ndarray = None) -> Tuple[np.ndarray, str, int]:
	"""Project a single reference trajectory onto the feasible set.

	Args:
	X_ref: (dof,) flat DOF vector in normalized space.
	s0: (state_dim,) initial state in normalized space.
	reg_weight: optional regularization weight (overrides default 1.0).
	a_prev: (action_dim,) last executed action in normalized space (for action rate constraint).

	Returns:
	X_proj: (dof,) projected DOF vector.
	status: "optimal" if the solver converged, "failed" otherwise.
	n_iters: number of IPOPT iterations spent on the solve.
	"""
	X_ref = np.asarray(X_ref).flatten()
	s0 = np.asarray(s0).flatten()

	assert X_ref.shape == (self.dof,), f"X_ref shape {X_ref.shape} != ({self.dof},)"
	assert s0.shape == (self.state_dim,), f"s0 shape {s0.shape} != ({self.state_dim},)"

	# CasADi parameters are column vectors
	self._opti.set_value(self._X_ref, X_ref.reshape(-1, 1))
	self._opti.set_value(self._s0, s0.reshape(-1, 1))
	if reg_weight is not None:
	self._opti.set_value(self._reg_weight, reg_weight)
	if a_prev is not None and self._a_prev is not None:
	self._opti.set_value(self._a_prev, np.asarray(a_prev).reshape(-1, 1))
	self._opti.set_initial(self._X, X_ref.reshape(-1, 1))

	try:
	sol = self._opti.solve()
	X_proj = np.array(sol.value(self._X)).flatten()
	status = "optimal"
	stats_src = self._opti
	except RuntimeError:
	X_proj = np.array(self._opti.debug.value(self._X)).flatten()
	status = "failed"
	stats_src = self._opti.debug

	try:
	n_iters = int(stats_src.stats()["iter_count"])
	except Exception:
	n_iters = 0

	return X_proj, status, n_iters

	def project_batch(
	self,
	X_refs: np.ndarray,
	s0s: np.ndarray,
	n_workers: int = 1,
	reg_weight: float = None,
	a_prevs: np.ndarray = None,
	) -> Tuple[np.ndarray, List[str], List[int]]:
	"""Project a batch of trajectories in parallel.

	Args:
	X_refs: (batch, dof)
	s0s: (batch, state_dim)
	n_workers: number of parallel IPOPT solves.
	reg_weight: optional regularization weight (passed to each solve).
	a_prevs: (batch, action_dim) last executed actions for action rate constraint.

	Returns:
	X_projs: (batch, dof)
	statuses: list of per-solve statuses ("optimal"/"failed").
	iters: list of per-solve IPOPT iteration counts.
	"""
	batch_size = X_refs.shape[0]
	assert X_refs.shape == (batch_size, self.dof)
	assert s0s.shape == (batch_size, self.state_dim)

	if n_workers <= 1:
	projs = []
	statuses = []
	iters = []
	for i in range(batch_size):
	ap = a_prevs[i] if a_prevs is not None else None
	x_proj, status, n_iters = self.project(X_refs[i], s0s[i], reg_weight=reg_weight, a_prev=ap)
	projs.append(x_proj)
	statuses.append(status)
	iters.append(n_iters)
	return np.stack(projs), statuses, iters

	# parallel: build config dict for picklable worker function
	config = {
	"circular_obstacles": self.circular_obstacles,
	"planar_obstacles": self.planar_obstacles,
	"horizon": self.horizon,
	"action_dim": self.action_dim,
	"state_dim": self.state_dim,
	"normalizer": self.normalizer,
	"dynamics": self.dynamics,
	"dynamics_relaxation": self.dynamics_relaxation,
	"obstacle_margin": self.obstacle_margin,
	"input_saturation": self.input_saturation,
	"objective_fn": self.objective_fn,
	"objective_scale": self.objective_scale,
	"max_action_delta": self.max_action_delta,
	"solver_opts": self._solver_opts,
	}

	results = [None] * batch_size
	statuses = [None] * batch_size
	iters = [0] * batch_size
	with ProcessPoolExecutor(max_workers=n_workers) as pool:
	futures = {
	pool.submit(
	_build_and_solve, config, X_refs[i], s0s[i], reg_weight,
	a_prevs[i] if a_prevs is not None else None,
	): i
	for i in range(batch_size)
	}
	for future in as_completed(futures):
	idx = futures[future]
	x_proj, status, n_iters = future.result()
	results[idx] = x_proj
	statuses[idx] = status
	iters[idx] = n_iters

	return np.stack(results), statuses, iters

	# ------------------------------------------------------------------
	# Utilities
	# ------------------------------------------------------------------

	def dof_to_full_trajectory(self, X_dof: np.ndarray, s0: np.ndarray) -> np.ndarray:
	"""Reconstruct the full (horizon, transition_dim) trajectory from DOF + s0.

	Returns:
	traj: (horizon, action_dim + state_dim)
	"""
	X_dof = np.asarray(X_dof).flatten()
	s0 = np.asarray(s0).flatten()
	ad = self.action_dim
	sd = self.state_dim

	traj = np.zeros((self.horizon, ad + sd))

	# timestep 0: action from DOF, state = s0
	a0_s, a0_e = self._action_indices[0]
	traj[0, :ad] = X_dof[a0_s:a0_e]
	traj[0, ad:] = s0

	# timesteps 1..H-1: action + state from DOF
	for t in range(1, self.horizon):
	a_s, a_e = self._action_indices[t]
	s_s, s_e = self._state_indices[t - 1]
	traj[t, :ad] = X_dof[a_s:a_e]
	traj[t, ad:] = X_dof[s_s:s_e]

	return traj

	def full_trajectory_to_dof(self, traj: np.ndarray) -> np.ndarray:
	"""Extract the DOF vector from a full (horizon, transition_dim) trajectory.

	s_0 is dropped (it's a parameter, not a DOF).
	"""
	ad = self.action_dim
	sd = self.state_dim
	X_dof = np.zeros(self.dof)

	# a_0
	a0_s, a0_e = self._action_indices[0]
	X_dof[a0_s:a0_e] = traj[0, :ad]

	# a_t, s_t for t >= 1
	for t in range(1, self.horizon):
	a_s, a_e = self._action_indices[t]
	s_s, s_e = self._state_indices[t - 1]
	X_dof[a_s:a_e] = traj[t, :ad]
	X_dof[s_s:s_e] = traj[t, ad:]

	return X_dof