arXiv:2511.18157cs.LGcs.RO2025-11中稿 · the 1st Workshop o…被引 3

让3D空间变换在任意深度学习框架中可微计算,支持GPU加速与自动微分。

scipy.spatial.transform: Differentiable Framework-Agnostic 3D Transformations in Python

  • 重构SciPy的3D变换模块,兼容JAX/PyTorch/CuPy等数组库。
  • 支持向量化批量处理、GPU/TPU执行及原生自动微分,性能显著提升。
  • 适用于机器人、视觉、物理模拟中的可微建模,适合科研与工业级应用。

三维刚体变换(旋转与平移)是现代机器人学、计算机视觉与仿真中可微机器学习流程的核心。然而,由于坐标轴约定、归一化、组合一致性等问题,特别是SO(3)上的实现极易出错,且仅在边缘情况显现。旧版SciPy.spatial.transform模块虽经严格测试,但仅支持NumPy,限制了其在基于GPU加速和自动微分的工作流中的应用。本文对SciPy.spatial.transform功能进行全面重构,使其兼容任何实现Python数组API的数组库(如JAX、PyTorch、CuPy)。新实现保留原有接口,同时支持GPU/TPU执行、JIT编译、向量化批量处理以及通过后端原生自动微分。我们通过两个案例验证:(i) 3D变换与旋转的可扩展性;(ii) 基于JAX的无人机仿真,利用SciPy的Rotation实现旋转动力学的精确积分。相关贡献已合并至SciPy主干,将在下一版本发布,为可微系统与机器学习提供框架无关、生产级别的3D空间数学基础。

原文摘要 · Abstract (English)

Three-dimensional rigid-body transforms, i.e. rotations and translations, are central to modern differentiable machine learning pipelines in robotics, vision, and simulation. However, numerically robust and mathematically correct implementations, particularly on SO(3), are error-prone due to issues such as axis conventions, normalizations, composition consistency and subtle errors that only appear in edge cases. SciPy's spatial$.$transform module is a rigorously tested Python implementation. However, it historically only supported NumPy, limiting adoption in GPU-accelerated and autodiff-based workflows. We present a complete overhaul of SciPy's spatial$.$transform functionality that makes it compatible with any array library implementing the Python array API, including JAX, PyTorch, and CuPy. The revised implementation preserves the established SciPy interface while enabling GPU/TPU execution, JIT compilation, vectorized batching, and differentiation via native autodiff of the chosen backend. We demonstrate how this foundation supports differentiable scientific computing through two case studies: (i) scalability of 3D transforms and rotations and (ii) a JAX drone simulation that leverages SciPy's Rotation for accurate integration of rotational dynamics. Our contributions have been merged into SciPy main and will ship in the next release, providing a framework-agnostic, production-grade basis for 3D spatial math in differentiable systems and ML.

3D变换可微计算JAXPyTorch

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。