如何解决Drake中AutoDiffXd与np.linalg.solve的兼容问题

问题描述 #

我需要求解形如M(q, dq)*ddq + C(q, dq) + G(q) = f的动力学方程以得到ddq，但遇到np.linalg.solve（及scipy变体）与Drake的AutoDiffXd类型不兼容的错误，该如何解决？

使用环境 #

• Ubuntu 24.04
• Drake 1.34.0
• numpy 2.1.3
• Python 3.12虚拟环境

原始代码 #

from pydrake.systems.framework import LeafSystem_, Context, VectorBase, ContinuousState
from pydrake.systems.scalar_conversion import TemplateSystem

import numpy as np

from BallbotMPC.ballbot_params import BallbotParams



@TemplateSystem.define("BallbotSystem_")
def BallbotSystem_(T):
    class Impl(LeafSystem_[T]):
        def _construct(self, converter=None, params=BallbotParams()):
            LeafSystem_[T].__init__(self, converter=converter)

            self._params = params

            self.DeclareContinuousState(4) # q: 4 states [ball angle, ball velocity, lean angle, lean velocity]

            self._command_port = self.DeclareVectorInputPort("u", 4)

            self._state_port = self.DeclareVectorOutputPort("state", 4, self.CopyStateOut)


        def _construct_copy(self, other, converter=None):
            Impl._construct(self, converter=converter, params=other._params)


        def DoCalcTimeDerivatives(self, context: Context, derivatives: ContinuousState):
            # constants
            mass = self._params.mass_k + self._params.mass_a + self._params.mass_w
            radius_total = self._params.radius_k + self._params.radius_w
            gam = self._params.l * self._params.mass_a + radius_total * self._params.mass_w

            state = self._state_port.Eval(context)
            u = self._command_port.Eval(context)

            # Dtype, necessary for numpy to understand the type of the drake types (AutoDiffXd, etc.)
            dtype = type(state[0])
            print(dtype)

            # Mass term in Langrangian formulation (Eqn. 2.22)
            M_x = np.zeros((2, 2), dtype=dtype)
            M_x[0, 0] = mass * self._params.radius_k ** 2 + self._params.Theta_k + (self._params.radius_k / self._params.radius_w) ** 2 * self._params.Theta_w
            M_x[0, 1] = -(self._params.radius_k / self._params.radius_w ** 2) * radius_total * self._params.Theta_w + gam * self._params.radius_k * np.cos(state[2])
            M_x[1, 0] = M_x[0, 1]
            M_x[1, 1] = (radius_total / self._params.radius_w) ** 2 * self._params.Theta_w + self._params.Theta_a + self._params.mass_a * self._params.l ** 2 + self._params.mass_w * radius_total ** 2

            # Coriolis term in Langrangian formulation (Eqn. 2.23)
            C_x = np.zeros(2, dtype=dtype)
            C_x[0] = -self._params.radius_k * gam * np.sin(state[2]) * state[3] ** 2

            # Gravitational term in Lagrangian formulation (Eqn. 2.24)
            G_x = np.zeros(2, dtype=dtype)
            G_x[1] = -self._params.g * np.sin(state[2]) * gam

            # Non-potential force term in Lagrangian formulation (Eqn. 2.17 + Eqn. 2.18)
            f_np = np.zeros(2, dtype=dtype)
            f_np[0] = self._params.radius_k / self._params.radius_w * u
            f_np[1] = -f_np[0]

            # Solve for ddq in M(q, dq)*ddq + C(q, dq) + G(q) = f_np
            ddq = np.linalg.solve(M_x, f_np - C_x - G_x)


            derivatives.get_mutable_vector().SetAtIndex(0, state[1])
            derivatives.get_mutable_vector().SetAtIndex(1, ddq[0])
            derivatives.get_mutable_vector().SetAtIndex(2, state[3])
            derivatives.get_mutable_vector().SetAtIndex(3, ddq[1])


        def CopyStateOut(self, context: Context, output: VectorBase):
            state = context.get_continuous_state_vector().CopyToVector()
            output.SetFromVector(state)
            

    return Impl


BallbotSystem = BallbotSystem_[None]

解决方案 #

numpy的np.linalg.solve不支持Drake的AutoDiffXd自定义自动微分类型，需改用Drake内置的线性求解工具eigen_solve，它基于Eigen实现，完美兼容AutoDiffXd并能保留导数信息。

修改步骤： #

1. 导入Drake的eigen_solve工具及Eigen矩阵/向量类型：

from pydrake.common import eigen_solve
from pydrake.math import Matrix2, Vector2

2. 替换numpy矩阵/向量的构建方式为Drake的Eigen类型：
   在DoCalcTimeDerivatives函数中，将所有矩阵和向量的定义改为Eigen类型：

# 构建质量矩阵
M_x = Matrix2[dtype]()
M_x[0, 0] = mass * self._params.radius_k ** 2 + self._params.Theta_k + (self._params.radius_k / self._params.radius_w) ** 2 * self._params.Theta_w
M_x[0, 1] = -(self._params.radius_k / self._params.radius_w ** 2) * radius_total * self._params.Theta_w + gam * self._params.radius_k * np.cos(state[2])
M_x[1, 0] = M_x[0, 1]
M_x[1, 1] = (radius_total / self._params.radius_w) ** 2 * self._params.Theta_w + self._params.Theta_a + self._params.mass_a * self._params.l ** 2 + self._params.mass_w * radius_total ** 2

# 构建Coriolis项
C_x = Vector2[dtype]()
C_x[0] = -self._params.radius_k * gam * np.sin(state[2]) * state[3] ** 2
C_x[1] = 0.0

# 构建重力项
G_x = Vector2[dtype]()
G_x[0] = 0.0
G_x[1] = -self._params.g * np.sin(state[2]) * gam

# 构建非保守力项
f_np = Vector2[dtype]()
f_np[0] = self._params.radius_k / self._params.radius_w * u
f_np[1] = -f_np[0]

3. 使用eigen_solve求解线性方程：

# 计算右侧向量
b = f_np - C_x - G_x
# 求解ddq
ddq = eigen_solve(M_x, b)

完整修改后的关键代码片段： #

def DoCalcTimeDerivatives(self, context: Context, derivatives: ContinuousState):
    # constants
    mass = self._params.mass_k + self._params.mass_a + self._params.mass_w
    radius_total = self._params.radius_k + self._params.radius_w
    gam = self._params.l * self._params.mass_a + radius_total * self._params.mass_w

    state = self._state_port.Eval(context)
    u = self._command_port.Eval(context)

    dtype = type(state[0])

    # 使用Eigen类型构建矩阵和向量
    M_x = Matrix2[dtype]()
    M_x[0, 0] = mass * self._params.radius_k ** 2 + self._params.Theta_k + (self._params.radius_k / self._params.radius_w) ** 2 * self._params.Theta_w
    M_x[0, 1] = -(self._params.radius_k / self._params.radius_w ** 2) * radius_total * self._params.Theta_w + gam * self._params.radius_k * np.cos(state[2])
    M_x[1, 0] = M_x[0, 1]
    M_x[1, 1] = (radius_total / self._params.radius_w) ** 2 * self._params.Theta_w + self._params.Theta_a + self._params.mass_a * self._params.l ** 2 + self._params.mass_w * radius_total ** 2

    C_x = Vector2[dtype]()
    C_x[0] = -self._params.radius_k * gam * np.sin(state[2]) * state[3] ** 2
    C_x[1] = 0.0

    G_x = Vector2[dtype]()
    G_x[0] = 0.0
    G_x[1] = -self._params.g * np.sin(state[2]) * gam

    f_np = Vector2[dtype]()
    f_np[0] = self._params.radius_k / self._params.radius_w * u
    f_np[1] = -f_np[0]

    # 求解线性方程
    b = f_np - C_x - G_x
    ddq = eigen_solve(M_x, b)

    derivatives.get_mutable_vector().SetAtIndex(0, state[1])
    derivatives.get_mutable_vector().SetAtIndex(1, ddq[0])
    derivatives.get_mutable_vector().SetAtIndex(2, state[3])
    derivatives.get_mutable_vector().SetAtIndex(3, ddq[1])

原理说明 #

Drake的AutoDiffXd是为自动微分场景设计的类型，numpy的线性代数工具仅支持基础数值类型（如float、double），无法识别AutoDiffXd的微分信息。而eigen_solve基于Eigen库实现，完全兼容Drake的自动微分类型，能在求解线性方程的同时保留导数链，适合MPC、最优控制等需要自动微分的动力学场景。

内容的提问来源于stack exchange，提问作者Jacob Sullivan