ACADOS中非线性成本函数的高效实现与配置

import casadi as ca
import numpy as np
from acados_template import AcadosModel, AcadosOcp, AcadosOcpOptions

def mobile_robot_model():
    """
    定义一个简单的移动机器人模型。

    返回:
        model (AcadosModel): 包含机器人动力学信息的Acados模型对象。
    """

    model_name = 'mobile_robot'

    # 定义符号变量 (状态)
    x = ca.MX.sym('x')        # x坐标
    y = ca.MX.sym('y')        # y坐标
    v = ca.MX.sym('v')        # 速度
    theta = ca.MX.sym('theta') # 朝向角

    # 定义控制输入
    a = ca.MX.sym('a')  # 加速度
    w = ca.MX.sym('w')  # 角速度

    # 定义状态和控制向量
    states = ca.vertcat(x, y, v, theta)
    controls = ca.vertcat(a, w)

    # 定义连续时间动力学 (x_dot = f(x, u))
    rhs = [v * ca.cos(theta), v * ca.sin(theta), a, w]
    x_dot = ca.MX.sym('x_dot', len(rhs)) # 状态导数的符号变量

    # 创建CasADi函数表示连续时间动力学
    continuous_dynamics = ca.Function(
        'continuous_dynamics',
        [states, controls],
        [ca.vcat(rhs)],
        ["state", "control_input"],
        ["rhs"]
    )

    # 显式和隐式动力学表达式
    f_expl_expr = continuous_dynamics(states, controls)
    f_impl_expr = x_dot - f_expl_expr

    # 创建AcadosModel对象
    model = AcadosModel()
    model.f_expl_expr = f_expl_expr
    model.f_impl_expr = f_impl_expr
    model.x = states
    model.xdot = x_dot
    model.u = controls
    model.p = [] # 无额外参数
    model.name = model_name

    return model

def create_ocp_solver():
    # 创建AcadosOcp对象
    ocp = AcadosOcp()

    # 设置优化问题模型
    model = mobile_robot_model()
    ocp.model = model

    # --------------------参数设置--------------------
    nx = model.x.size()[0]  # 状态维度
    nu = model.u.size()[0]  # 控制输入维度
    N = 100                 # 离散化步数
    T = 30.0                # 总时间

    ocp.dims.N = N
    ocp.dims.nx = nx
    ocp.dims.nu = nu
    ocp.solver_options.tf = T

    # 初始状态
    x_ref = np.array([0.0, 0.0, 0.0, 0.0]) # 示例初始状态 [x, y, v, theta]
    ocp.constraints.x0 = x_ref

    # 参数初始化 (如果模型有参数，此处设置)
    ocp.dims.np = len(model.p)
    ocp.parameter_values = np.zeros(ocp.dims.np)

    # ---------------------约束设置---------------------
    # 控制输入约束
    ocp.constraints.lbu = np.array([-0.1, -0.3])  # 控制输入的下界 [a_min, w_min]
    ocp.constraints.ubu = np.array([0.1, 0.3])    # 控制输入的上界 [a_max, w_max]
    ocp.constraints.idxbu = np.array([0, 1])      # 对应控制输入向量的索引

    # 状态约束 (示例，实际根据需求设置)
    # ocp.constraints.lbx = np.array([-100, -100, 0, -np.pi]) # 状态下界
    # ocp.constraints.ubx = np.array([100, 100, 10, np.pi])   # 状态上界
    # ocp.constraints.idxbx = np.array([0, 1, 2, 3]) # 对应状态向量的索引

    return ocp

def setup_nonlinear_ls_cost(ocp):
    nx = ocp.dims.nx
    nu = ocp.dims.nu

    # 设置成本类型为非线性最小二乘
    ocp.cost.cost_type = 'NONLINEAR_LS'
    ocp.cost.cost_type_e = 'NONLINEAR_LS' # 终端成本也使用非线性最小二乘

    # 1. 定义输出表达式 y(x, u)
    # 假设我们想跟踪 x, y, theta 和控制输入 a, w
    # 那么 y_expr 应该包含这些项
    q_expr = ca.vertcat(ocp.model.x[0], ocp.model.x[1], ocp.model.x[3]) # x, y, theta
    u_expr = ocp.model.u # a, w
    ocp.model.cost_y_expr = ca.vertcat(q_expr, u_expr) # 中间阶段输出
    ocp.model.cost_y_expr_e = q_expr # 终端阶段输出 (通常只关心状态)

    # 2. 定义权重矩阵 W
    # 假设 q 的维度是 3 (x,y,theta)，u 的维度是 2 (a,w)
    ny = ocp.model.cost_y_expr.size()[0] # 中间阶段输出维度 (3+2=5)
    ny_e = ocp.model.cost_y_expr_e.size()[0] # 终端阶段输出维度 (3)

    # 权重矩阵 W (对角矩阵)
    # 对应 [x_error, y_error, theta_error, a_error, w_error]
    W_diag = np.array([100.0, 100.0, 50.0, 1.0, 1.0])
    ocp.cost.W = np.diag(W_diag)

    # 终端权重矩阵 W_e (对角矩阵)
    # 对应 [x_error, y_error, theta_error]
    W_e_diag = np.array([200.0, 200.0, 100.0])
    ocp.cost.W_e = np.diag(W_e_diag)

    # 3. 定义参考值 y_ref
    # 这些是期望的输出值，可以在运行时更新
    # 示例：假设期望轨迹是原点，期望控制输入是零
    ocp.cost.yref = np.zeros(ny) # 中间阶段参考值
    ocp.cost.yref_e = np.zeros(ny_e) # 终端阶段参考值

    # 可以通过 ocp.cost.set(stage, "yref", new_yref_value) 在求解前更新 yref
    # 例如：ocp.cost.set(0, "yref", np.array([1.0, 1.0, 0.0, 0.0, 0.0]))

    return ocp

def setup_external_cost(ocp):
    # 设置成本类型为外部成本
    ocp.cost.cost_type = 'EXTERNAL'
    ocp.cost.cost_type_e = 'EXTERNAL' # 终端成本也使用外部成本

    # 假设障碍物位置为 (ox, oy)，半径为 r_obstacle
    # 机器人半径为 r_robot
    ox = 5.0
    oy = 5.0
    r_obstacle = 0.5
    r_robot = 0.2

    # 机器人当前位置 (x, y)
    robot_x = ocp.model.x[0]
    robot_y = ocp.model.x[1]

    # 计算机器人与障碍物中心的距离
    distance_sq = (robot_x - ox)**2 + (robot_y - oy)**2
    distance = ca.sqrt(distance_sq)

    # 安全距离裕度
    safety_margin = r_robot + r_obstacle

    # 定义避障成本表达式
    # 当距离小于安全距离时，成本急剧增加
    # 使用 max(0, safety_margin - distance) 避免负距离倒数
    # 或者更直接地，使用一个平滑的惩罚函数

    # 方式一：直接倒数惩罚 (可能导致数值问题，需要截断或平滑)
    # cost_avoidance = 0.5 * (1 / (distance - safety_margin + 1e-6))**2
    # cost_avoidance = ca.if_else(distance < safety_margin, 1e6, 0.0) # 简单粗暴的惩罚

    # 方式二：更平滑的惩罚函数，避免数值问题
    # 例如，当距离小于某个阈值时，成本开始增加
    threshold_distance = safety_margin + 0.5 # 在安全距离外0.5米开始惩罚

    # 定义一个惩罚项，当距离小于阈值时，惩罚值非零
    # 惩罚函数 f(d) = max(0, threshold_distance - d)^2
    # 权重系数 beta 调整惩罚强度
    beta = 1000.0 # 避障权重

    # 计算距离差值，如果为负，则说明在安全距离内
    dist_diff = threshold_distance - distance

    # 只有当 dist_diff > 0 (即距离小于阈值) 时才产生惩罚
    # 使用 CasADi 的 if_else 来实现条件逻辑
    cost_avoidance = beta * ca.if_else(dist_diff > 0, dist_diff**2, 0.0)

    ocp.model.cost_expr_ext_cost = cost_avoidance # 中间阶段避障成本
    ocp.model.cost_expr_ext_cost_e = 0.0 # 终端阶段通常不考虑避障，或设置为0

    return ocp