Blame for applications/avoiding/avoiding_baselines/DiRecT/projector.py

"""
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