leap_c.diff_mpc.data

Classes

AcadosOcpSolverInput

Input for an AcadosOcpSolver representing a batch of problem instances to be solved.

Functions

collate_acados_flattened_batch_iterate_fn(...)
collate_acados_flattened_iterate_fn(...)
collate_acados_ocp_solver_input(→ AcadosOcpSolverInput)	Collates a batch of AcadosOcpSolverInput objects into a single object.
validate_forward_inputs(→ None)	Validate the shapes of x0 and u0 before the forward solve.

Module Contents

class leap_c.diff_mpc.data.AcadosOcpSolverInput

Bases: NamedTuple

Input for an AcadosOcpSolver representing a batch of problem instances to be solved.

Variables:

• x0 – Initial state, shape (batch_size, nx)

• u0 – Initial control input, shape (batch_size, nu), optional. If provided, the initial control input will be constrained to this.

• p_global – Global parameters, shape (batch_size, np_global), optional. If not provided, the default values set in the acados ocp object will be used.

• p_stagewise – Stage-wise parameters, shape (batch_size, N_horizon + 1, np_stagewise), or (batch_size, len(p_stagewise_sparse_idx), np_stagewise), if p_stagewise_sparse_idx is provided, optional. If not provided, the default values set in the acados ocp object will be used. Has to be provided if p_stagewise_sparse_idx is provided.

• p_stagewise_sparse_idx – If provided, the indices determine which elements of the total stagewise parameter in the solver should be overwritten by the provided p_stagewise values, shape (batch_size, N_horizon + 1, nindices), optional.

get_sample(idx: int) → AcadosOcpSolverInput

Get the sample at index idx from the batch.

property batch_size: int

Get the batch size.

p_global: numpy.ndarray | None = None

p_stagewise: numpy.ndarray | None = None

p_stagewise_sparse_idx: numpy.ndarray | None = None

u0: numpy.ndarray | None = None

x0: numpy.ndarray

leap_c.diff_mpc.data.collate_acados_flattened_batch_iterate_fn(batch: Sequence[acados_template.acados_ocp_iterate.AcadosOcpFlattenedBatchIterate], collate_fn_map: dict | None = None) → acados_template.acados_ocp_iterate.AcadosOcpFlattenedBatchIterate

leap_c.diff_mpc.data.collate_acados_flattened_iterate_fn(batch: Sequence[acados_template.acados_ocp_iterate.AcadosOcpFlattenedIterate], collate_fn_map: dict | None = None) → acados_template.acados_ocp_iterate.AcadosOcpFlattenedBatchIterate

leap_c.diff_mpc.data.collate_acados_ocp_solver_input(batch: Sequence[AcadosOcpSolverInput], collate_fn_map: dict | None = None) → AcadosOcpSolverInput

Collates a batch of AcadosOcpSolverInput objects into a single object.

leap_c.diff_mpc.data.validate_forward_inputs(ocp: acados_template.AcadosOcp, x0: numpy.ndarray | torch.Tensor, u0: numpy.ndarray | torch.Tensor | None) → None

Validate the shapes of x0 and u0 before the forward solve.

Works on any array-like exposing ndim and shape (numpy, torch, jax), so it can be shared by the torch and jax layers. Run before solver initialization so dimension mismatches fail fast with a clear message rather than as a cryptic C-level error.

Raises:

ValueError – If x0 is not 2-D, its last dimension is not ocp.dims.nx, u0 is not 2-D, its batch does not match x0, or its last dimension is not ocp.dims.nu.