跳到主要内容

开发自己的 LambdaSim Solver

这篇文档面向希望为 LambdaSim 接入新求解器的开发者。你将先实现一个可以运行、重置和释放的最小 Solver，再学习如何读取场景数据、发布计算结果，以及参考 IPC 和 MuJoCo 接入真实物理算法。

Solver 的接入流程是：读取 Scene 输入 → 构建求解器 → 接收时间步 → 发布结果到 SimData。 LambdaSim 负责组织场景和调用顺序，具体的物理算法、内部状态及资源由 Solver 管理。

Solver 基础概念​

Solver 基类​

SolverBase 是 LambdaSim 求解器的统一基类。开发者通过继承它，将自己的计算算法接入 Scene。

Solver 注册到 Scene 后，由 Scene 管理其构建、推进、重置和释放。具体算法、内部状态和计算资源由 Solver 自己实现和管理。

from lambdasim.solvers import SolverBase

源码位置：src/lambdasim/solvers/soloverbase.py。

成员方法​

接口职责
finalize__ektro(scene, num_envs, sim_data)从 Scene 获取输入，构建模型、初始状态和计算资源，建立公共输出。
step__ektro(step_context)按传入的时间步执行计算，并写回公共结果。
on_reset()恢复构建时的初始状态和公共输出。
close()释放 Solver 持有的资源。
upload_control__ektro(sim_data)显式上传控制目标，由应用按需调用。
from_json(config)解析该 Solver 支持的配置，由应用按需调用。

实现新的 Solver 时，首先实现 finalize__ektro() 和 step__ektro()；如果持有可变状态或计算资源，还需要实现相应的on_reset()和close()逻辑。

成员变量​

基类提供以下配置成员：

成员用途
src_paths指定参与计算的场景路径。
src_tags指定参与计算的语义标签。
src_config描述输入选择的层级范围，以及正选或反选方式。
substep_count描述 Solver 内部的子步数量。

这些成员用于表达配置，基类不会自动执行输入筛选或内部子步计算。具体 Solver 需要按照自己的契约使用这些配置。

向 Scene 添加 Solver​

以 MuJoCo Ektro 求解器为例，通过 add_solver__ektro() 将实例添加到 Scene：

from lambdasim.scene.scene import Scene
from lambdasim.solvers.solver_mujoco_ektro import MujocoEktroSolver

scene = Scene()
solver = MujocoEktroSolver()
scene.add_solver__ektro(solver)

添加后，Solver 的构建、推进、重置和释放由 Scene 统一管理。这里仅示意注册过程，实际计算前还需要准备输入并完成初始化。

Solver 计算更新​

每次循环推进一个Step，由 Scene 按注册顺序调用各个 Solver 执行计算。以下假设场景已完成初始化：

from lambdasim.scene.step_context import StepContext

context = StepContext(step_index=0, time=0.0, dt=1.0 / 60.0)

while True:
    # 推进一帧的目标时间并执行计算。Advance the target time by one frame and evaluate.
    context.set_time(context.time + context.dt)
    scene.step__ektro(context)

这里每帧推进 1/60 秒仿真时间；若求解器需要内部子步，由其自身按配置执行。

自定义 Solver​

继承 SolverBase，实现自己的求解器。Scene 构建时会调用 Solver 的 finalize__ektro() 方法，将场景、环境数量和公共运行数据传入。

这里假设调用方已通过 solver.src_paths 配置至少一个输入路径，示例只展示读取第一个输入。具体需要哪些数据，由求解器自身的计算需求决定。

from lambdasim.solvers import SolverBase


class YOUR_SOLVER_NAME(SolverBase):
    def finalize__ektro(self, scene, num_envs, sim_data):
        path = self.src_paths[0]

        # 从 Scene 读取几何、物性和变换。Read geometry, physics and transforms from Scene.
        mesh = scene.get_mesh_data(path, space="local")
        physics = scene.get_physics_data(path)
        world = scene.get_world_xform(path)

        # 保存公共运行数据引用。Keep the shared runtime data reference.
        self.sim_data = sim_data

构建阶段的输入数据​

在上面的 finalize__ektro() 示例中，scene、num_envs 和 sim_data 由 Scene 调用时传入；mesh、physics 和 world 则是 Solver 根据输入路径，从 Scene 查询得到的数据。

名称说明
mesh几何点位置及拓扑等数据，示例中使用局部坐标空间。
physics场景中已有的物性和约束数据。
world输入对象的世界变换，用于处理局部坐标与世界坐标的关系。
num_envs本次构建的仿真环境数量。
sim_dataScene 管理的公共运行数据，计算设备由 sim_data.device 指定。

Solver 根据这些输入，在构建阶段完成数据校验、模型转换和资源分配，为后续 step__ektro() 的计算做好准备。

取得输入后，在 finalize__ektro() 中构建求解状态并建立公共输出；后续在 step__ektro() 中执行计算，将需要共享的结果写回 self.sim_data.data。算法内部的状态和工作区由 Solver 自己持有。

2. 实现第一个 Solver​

先实现一个匀速运动的一维状态：速度为 v（m/s），位置单位为米，每个时间步执行 position += v * dt。这个例子只用于验证接入协议，输出是 Python 标量，不加载资产、不使用 GPU，也不更新视口。

将下面的类和下一节的运行代码放在同一个脚本中，即可从仓库根目录运行。临时练习脚本可以放在 tmp/。

from lambdasim.solvers import SolverBase
from lambdasim.scene.scene import Scene
from lambdasim.scene.step_context import StepContext

class ConstantVelocitySolver(SolverBase):
    """最小的一维状态更新。A minimal one-dimensional state update."""

    def __init__(self, velocity: float = 1.0) -> None:
        super().__init__()
        self.velocity = velocity
        self._output = None

    def finalize__ektro(self, scene, num_envs, sim_data) -> None:
        self._output = {"position": 0.0}
        sim_data.data["tutorial_motion"] = self._output

    def step__ektro(self, step_context: StepContext) -> None:
        self._output["position"] += self.velocity * step_context.dt

这个类分别完成了四件事：

方法本例做什么接入真实算法时做什么
__init__()保存速度。保存配置；此时不分配设备状态、不打开场景资产。
finalize__ektro()校验环境数，创建公共输出组。读取并校验输入，构建模型和初态，分配工作区，发布初始输出。
step__ektro()用本次 dt 更新位置。执行一次 Scene 子步，随后更新公共输出。
on_reset() / close()归零位置 / 释放引用。恢复构建初态 / 释放持有的资源。

这里的 tutorial_motion 是教程自定义的数据组，不是 LambdaSim 预定义的物理 schema。接入大量粒子、顶点或刚体时，应按第 6 节定义 Ektro 数组及其语义。

3. 接入 Scene 并运行​

在上面的类定义后添加：

scene = Scene()
solver = scene.add_solver__ektro(ConstantVelocitySolver(velocity=2.0))

scene.finalize__ektro(device="cpu", num_envs=1)
output = scene.sim_data.data["tutorial_motion"]

request = StepContext(step_index=0, time=0.0, dt=1.0 / 60.0)
for step in range(1, 6):
    request.set_time(step * request.dt)
    scene.step__ektro(request)
    print(f"step = {step}, time = {request.time:.6f}, position = {output['position']:.6f}")

输出：

step = 1, time = 0.016667, position = 0.033333
step = 2, time = 0.033333, position = 0.066667
step = 3, time = 0.050000, position = 0.100000
step = 4, time = 0.066667, position = 0.133333
step = 5, time = 0.083333, position = 0.166667

你已经完成 Solver 的最小接入。Scene 会保存注册的实例，finalize__ektro() 向它传入本次构建的 SimData。后续调用方通过 scene.sim_data 访问同一份公共输出。

device 是必填参数，不依赖进程默认设备。正式使用 Vulkan 或 MUSA 时，应明确指定可用的逻辑设备，例如 musa:0，并确认底层算法支持它。

5. 获取场景中的物体​

Solver 通常在 finalize__ektro(scene, num_envs, sim_data) 中，从已经加载资产的 Scene 获取输入，并保存后续计算需要的数据引用。

Scene 提供以下几类数据访问入口

获取方式接口获取内容
按名称获取或遍历物体scene.get_object(name)
scene.object_list		场景物体及其几何、属性和位姿
按场景路径查询和读取scene.query_paths(...)
scene.get_*_data(path)
scene.get_world_xform(path)		网格、物性、碰撞体和变换等数据
获取结构化刚体描述scene.rigid_entity刚体、关节、碰撞体和机器人结构
获取仿真运行数据scene.sim_data.dataSolver 发布的位置、速度等计算结果

根据 Object Name 获取物体​

通过 scene.get_object(name) 获取场景物体：

from lambdasim.solvers import SolverBase


class MySolver(SolverBase):
    def finalize__ektro(self, scene, num_envs, sim_data):
        self.obj = scene.get_object("Deformable")

"Deformable" 是导入资产时指定的物体名称，返回值是 SceneObject。保存在 self.obj 后，可以在后续计算中访问该物体。

根据 TAG 标签查询路径并获取 GeoView​

先通过标签查询输入路径，再使用 scene.get_views() 获取对应的几何视图：

from lambdasim.solvers import SolverBase


class MySolver(SolverBase):
    def finalize__ektro(self, scene, num_envs, sim_data):
        self.mesh_paths = scene.query_paths(
            tags=["cloth"], # 多个标签之间是 “或（OR）” 的关系
            kind="mesh",
        )
        self.views = scene.get_views(self.mesh_paths)

self.mesh_paths 是路径列表，self.views 是与路径顺序对应的 GeoViewHandle 列表。通过视图可以直接读取几何属性：

排除指定 TAG 的几何输入​

使用 mode=SourceMode.EXCLUDE 反选。例如，获取场景中所有不带 collider 标签的几何视图：

from lambdasim.solvers import SolverBase, SourceMode


class MySolver(SolverBase):
    def finalize__ektro(self, scene, num_envs, sim_data):
        self.mesh_paths = scene.query_paths(
            tags=["collider"],
            kind="mesh",
            mode=SourceMode.EXCLUDE,
        )
        self.views = scene.get_views(self.mesh_paths)

从上级节点继承了 collider 标签的几何也会被排除。

获取指定层级下带 TAG 的几何​

例如，获取 /root/obj/Factory 层级内带 cloth 标签的几何：

from lambdasim.solvers import SolverBase, SourceScope


class MySolver(SolverBase):
    def finalize__ektro(self, scene, num_envs, sim_data):
        self.mesh_paths = scene.query_paths(
            paths=["/root/obj/Factory"],
            tags=["cloth"],
            kind="mesh",
            scope=SourceScope.DESCENDANTS,
        )
        self.views = scene.get_views(self.mesh_paths)

SourceScope.DESCENDANTS 查询指定节点自身及其全部后代。只有位于这个范围内、且匹配 cloth 标签的几何才会返回；标签可以从上级节点继承。

将 /root/obj/Factory 替换为实际场景路径。

6. 定义并发布公共输出​

SimData 不要求把算法所有数据都公开。模型、搜索结构、接触工作区和迭代暂存可以保留在 Solver 内；只公开需要交给其他模块的状态。

接入前，为每个公共数组明确：名称、业务域、dtype、shape、device、单位、坐标系、索引映射、读写归属。以下是现有实现的例子，不能跨类型直接替换：

公共字段当前存储语义
rigid.body_qfloat32[E, B, 7]世界位置与旋转，分量顺序为 px, py, pz, qx, qy, qz, qw。
rigid.body_qdfloat32[E, B, 6]世界空间线速度在前、角速度在后，即 vx, vy, vz, wx, wy, wz。
rigid.body_paths长度为 B 的路径序列与 body 数组索引对齐，供耦合消费者匹配实体。
ipc.<field>_positionsfloat32[1, 1, N, 3]IPC 某个可变形域的世界空间点位置；field 由 solver.output_fields[path] 查询。
ipc.<field>_velocitiesfloat32[1, 1, N, 3]对应可变形域的点速度。

E 是环境数，B 是刚体数，N 是该 IPC 域的点数。上述 IPC 布局是当前适配器布局，不是所有 Solver 的通用数组格式。Affine 域发布的是 surface_positions、q 和 qd，不能用布料的字段名读取。

公共数据遵守以下规则：

  • 一个组或字段明确由谁创建、谁写入，避免多个 Solver 覆盖同一状态。创建组前检查重名。
  • 消费者保留生产者的数组引用，通过路径映射匹配索引，不能假设两个 Solver 使用相同排序。
  • 状态更新和重置原地写入公共数组；改变拓扑和数组形状属于重新构建。
  • 计算设备来自 sim_data.device；资产数据需要上传时，构建阶段显式复制到目标设备。
  • 不支持某个 backend 或 num_envs 时，在构建阶段明确拒绝，不能只计算第一个环境或偷偷改用 CPU。

逐元素数值计算优先使用 Ektro 数学与 @ek.kernel。开始实现前先确定输入、输出、scratch 和每个 kernel 的运行域；转换、积分和结果发布可以是不同阶段。不要为了共享几个结果把私有求解工作区全部暴露到 SimData。

版权所有 2024 摩尔线程，本文档受国际版权法保护。

摩尔线程和摩尔线程徽标是摩尔线程智能科技 (北京)有限责任公司的注册商标。

免责声明

本文档提供有关摩尔线程产品的信息。本文档并未授权任何知识产权的许可，并未以明示或暗示，或以禁止反言或其他方式授予任何知识产权许可。