[Wave RHO][3/4] Preserve receding-horizon outcomes and expose Acados diagnostics

bioptim/__init__.py

@@ -246,6 +246,9 @@
     CyclicNonlinearModelPredictiveControl,
     CyclicMovingHorizonEstimator,
     MultiCyclicNonlinearModelPredictiveControl,
+    RecedingHorizonFailurePolicy,
+    RecedingHorizonWindowResult,
+    CyclicVariableShift,
 )
 from .optimization.receding_horizon_optimization import MovingHorizonEstimator, NonlinearModelPredictiveControl
 from .optimization.solution.solution import Solution

bioptim/interfaces/acados_interface.py

@@ -1,5 +1,6 @@
 from time import perf_counter
 from datetime import datetime
+from typing import TYPE_CHECKING
 
 import numpy as np
 from scipy import linalg
@@ -13,13 +14,24 @@
 from ..limits.path_conditions import Bounds
 from ..misc.enums import InterpolationType
 
+if TYPE_CHECKING:
+    from ..optimization.solution.solution import Solution
+
 
 from ..misc.parameters_types import (
     Str,
     Bool,
     AnyListorDict,
 )
 
+ACADOS_STATUS_LABELS = {
+    0: "success",
+    1: "maximum_iterations_reached",
+    2: "minimum_step_reached",
+    3: "qp_solver_failure",
+    4: "ready",
+}
+
 
 class AcadosInterface(SolverInterface):
     """
@@ -268,9 +280,19 @@ def __set_constraints(self, ocp) -> None:
             )
 
         for key in ocp.nlp[0].controls.keys():
-            if not np.all(np.all(ocp.nlp[0].u_bounds[key].min.T == ocp.nlp[0].u_bounds[key].min.T[0, :], axis=0)):
+            if not np.all(
+                np.all(
+                    ocp.nlp[0].u_bounds[key].min.T == ocp.nlp[0].u_bounds[key].min.T[0, :],
+                    axis=0,
+                )
+            ):
                 raise NotImplementedError("u_bounds min must be the same at each shooting point with ACADOS")
-            if not np.all(np.all(ocp.nlp[0].u_bounds[key].max.T == ocp.nlp[0].u_bounds[key].max.T[0, :], axis=0)):
+            if not np.all(
+                np.all(
+                    ocp.nlp[0].u_bounds[key].max.T == ocp.nlp[0].u_bounds[key].max.T[0, :],
+                    axis=0,
+                )
+            ):
                 raise NotImplementedError("u_bounds max must be the same at each shooting point with ACADOS")
 
             if (
@@ -478,7 +500,11 @@ def _adjust_dim():
                     v_var[rows] = 1.0
                     return v_var, rows
 
-                if objectives.node[0] not in [Node.INTERMEDIATES, Node.PENULTIMATE, Node.END]:
+                if objectives.node[0] not in [
+                    Node.INTERMEDIATES,
+                    Node.PENULTIMATE,
+                    Node.END,
+                ]:
                     v_var, rows = _adjust_dim()
                     if is_state:
                         acados.Vx0 = np.vstack((acados.Vx0, np.diag(v_var)))
@@ -517,7 +543,8 @@ def add_nonlinear_ls_lagrange(acados, objectives, t, dt, x, u, p, a, d):
                 u = vertcat(u, u)
 
             acados.lagrange_costs = vertcat(
-                acados.lagrange_costs, objectives.function[0](t, dt, x, u, p, a, d).reshape((-1, 1))
+                acados.lagrange_costs,
+                objectives.function[0](t, dt, x, u, p, a, d).reshape((-1, 1)),
             )
             acados.W = linalg.block_diag(
                 acados.W, np.diag(objectives.weight.evaluate_at(0, objectives.function[0].numel_out()))
@@ -530,7 +557,11 @@ def add_nonlinear_ls_lagrange(acados, objectives, t, dt, x, u, p, a, d):
                 acados.y_ref.append([np.zeros((objectives.function[0].numel_out(), 1)) for _ in node_idx])
 
         def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None):
-            if objectives.node[0] not in [Node.INTERMEDIATES, Node.PENULTIMATE, Node.END]:
+            if objectives.node[0] not in [
+                Node.INTERMEDIATES,
+                Node.PENULTIMATE,
+                Node.END,
+            ]:
                 acados.W_0 = linalg.block_diag(
                     acados.W_0, np.diag(objectives.weight.evaluate_at(0, objectives.function[0].numel_out()))
                 )
@@ -546,7 +577,8 @@ def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None):
                 u_tp = u_tp if objectives.function[0].size_in("u") != (0, 0) else []
 
                 acados.mayer_costs = vertcat(
-                    acados.mayer_costs, objectives.function[0](t, dt, x_tp, u_tp, p, a, d).reshape((-1, 1))
+                    acados.mayer_costs,
+                    objectives.function[0](t, dt, x_tp, u_tp, p, a, d).reshape((-1, 1)),
                 )
 
                 if objectives.target is not None:
@@ -569,7 +601,8 @@ def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None):
                 u_tp = u_tp if objectives.function[-1].size_in("u") != (0, 0) else []
 
                 acados.mayer_costs_e = vertcat(
-                    acados.mayer_costs_e, objectives.function[-1](t, dt, x_tp, u_tp, p, a, d).reshape((-1, 1))
+                    acados.mayer_costs_e,
+                    objectives.function[-1](t, dt, x_tp, u_tp, p, a, d).reshape((-1, 1)),
                 )
 
                 if objectives.target is not None:
@@ -590,7 +623,10 @@ def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None):
         self.W_e = np.zeros((0, 0))
         self.W_0 = np.zeros((0, 0))
         allowed_control_objectives = [ObjectiveFcn.Lagrange.MINIMIZE_CONTROL]
-        allowed_state_objectives = [ObjectiveFcn.Lagrange.MINIMIZE_STATE, ObjectiveFcn.Mayer.TRACK_STATE]
+        allowed_state_objectives = [
+            ObjectiveFcn.Lagrange.MINIMIZE_STATE,
+            ObjectiveFcn.Mayer.TRACK_STATE,
+        ]
 
         if self.acados_ocp.cost.cost_type == "LINEAR_LS":
             n_states = ocp.nlp[0].states.shape
@@ -875,6 +911,7 @@ def get_optimized_value(self) -> AnyListorDict:
             "iter": self.ocp_solver.get_stats("sqp_iter"),
             "status": self.status,
             "solver": SolverType.ACADOS,
+            "solver_diagnostics": self.get_solver_diagnostics(),
         }
 
         out["x"] = vertcat(out["x"], acados_x.reshape(-1, 1, order="F"))
@@ -893,6 +930,91 @@ def get_optimized_value(self) -> AnyListorDict:
 
         return out[0] if len(out) == 1 else out
 
+    def get_solver_diagnostics(self) -> dict:
+        """Return stable, best-effort diagnostics for the last Acados solve."""
+
+        diagnostics = {
+            "status": self.status,
+            "status_label": ACADOS_STATUS_LABELS.get(self.status, "unknown"),
+            "wall_time": self.real_time_to_optimize,
+        }
+        if self.ocp_solver is None:
+            return diagnostics
+
+        stat_names = (
+            "res_stat",
+            "res_eq",
+            "res_ineq",
+            "res_comp",
+            "sqp_iter",
+            "qp_iter",
+            "alpha",
+            "time_tot",
+            "statistics",
+        )
+        for name in stat_names:
+            try:
+                diagnostics[name] = self.ocp_solver.get_stats(name)
+            except Exception:
+                # The available statistics depend on the acados_template version and solver configuration.
+                continue
+        diagnostics["iterates"] = self.get_iterates()
+        return diagnostics
+
+    def get_iterates(self) -> list[dict]:
+        """Return the current Acados iterate without exposing acados_template internals."""
+
+        if self.ocp_solver is None:
+            raise RuntimeError("Acados must be solved before its iterates can be retrieved")
+        iterates = []
+        for node in range(self.acados_ocp.dims.N + 1):
+            iterate = {"x": np.asarray(self.ocp_solver.get(node, "x")).copy()}
+            for field in ("u", "z", "pi", "lam", "sl", "su"):
+                if field in ("u", "pi") and node == self.acados_ocp.dims.N:
+                    iterate[field] = None
+                    continue
+                try:
+                    iterate[field] = np.asarray(self.ocp_solver.get(node, field)).copy()
+                except Exception:
+                    iterate[field] = None
+            iterates.append(iterate)
+        return iterates
+
+    def set_iterates(self, iterates: list[dict], include_multipliers: bool = False) -> None:
+        """Initialize Acados from compatible iterates.
+
+        Primal fields (``x``, ``u`` and ``z``) are always transferred. Multipliers (``pi``, ``lam``, ``sl`` and
+        ``su``) are transferred only when explicitly requested; their dimensions must already match the Acados OCP.
+        """
+
+        if self.ocp_solver is None:
+            raise RuntimeError("The Acados solver must be created before its iterates can be initialized")
+        if len(iterates) != self.acados_ocp.dims.N + 1:
+            raise ValueError(f"Expected {self.acados_ocp.dims.N + 1} Acados nodes, got {len(iterates)}")
+
+        primal_fields = ("x", "u", "z")
+        multiplier_fields = ("pi", "lam", "sl", "su") if include_multipliers else ()
+        for node, iterate in enumerate(iterates):
+            for field in primal_fields + multiplier_fields:
+                value = iterate.get(field)
+                if value is None:
+                    continue
+                if field in ("u", "pi") and node == self.acados_ocp.dims.N:
+                    raise ValueError(f"Acados field '{field}' is not defined at the terminal node")
+                self.ocp_solver.set(node, field, np.asarray(value))
+
+    def set_lagrange_multiplier(self, sol: "Solution") -> None:
+        """Transfer Acados multipliers only from a compatible Acados iterate."""
+
+        diagnostics = getattr(sol, "solver_diagnostics", None) or {}
+        iterates = diagnostics.get("iterates")
+        if iterates is None:
+            raise ValueError(
+                "IPOPT/NLP multipliers cannot be converted generically to Acados QP multipliers. "
+                "Use a primal-only warm start, or provide compatible Acados iterates."
+            )
+        self.set_iterates(iterates, include_multipliers=True)
+
     def solve(self, expand_during_shake_tree: Bool = False) -> AnyListorDict:
         """
         Solve the prepared ocp

bioptim/optimization/optimal_control_program.py

@@ -1475,14 +1475,17 @@ def set_ocp_solver(self, solver: Solver) -> None:
                     "You cannot change the solver once it has been set. Please create a new OptimalControlProgram."
                 )
 
-    def set_warm_start(self, sol: Solution) -> None:
+    def set_warm_start(self, sol: Solution, transfer_multipliers: Bool | None = None) -> None:
         """
         Modify x and u initial guess based on a solution.
 
         Parameters
         ----------
         sol: Solution
             The solution to initiate the OCP from
+        transfer_multipliers: bool | None
+            Whether to transfer solver multipliers. By default they are reused only for IPOPT, whose NLP multiplier
+            layout is compatible. Acados QP multipliers require explicit opt-in and Acados-origin iterates.
         """
 
         state = sol.decision_states(to_merge=SolutionMerge.NODES)
@@ -1516,7 +1519,9 @@ def set_warm_start(self, sol: Solution) -> None:
 
         self.update_initial_guess(x_init=x_init_guess, u_init=u_init_guess, parameter_init=param_init_guess)
 
-        if self.ocp_solver:
+        if transfer_multipliers is None:
+            transfer_multipliers = self.ocp_solver is not None and self.ocp_solver.opts.type == SolverType.IPOPT
+        if self.ocp_solver and transfer_multipliers:
             self.ocp_solver.set_lagrange_multiplier(sol)
 
         self._is_warm_starting = True