Skip to content

perf: migrate MotrixSim write compiler to batched compile_write FFI plan #20

Description

@wlgys8

Summary

motrixsim(当前锁定 0.10.1.dev121620)已提供写侧批量 FFI 接口 SceneModel.compile_write(fields) -> WritePlan / WriteProgram(见 motrixsim/write.pyi),与读侧 compile_query 对应。读侧已迁移完成(motrix_env_motrixsim/src/motrix_env_motrixsim/sim_data.py:82),但写侧 write_compiler.py 仍是逐 op 的 Python 循环实现,未合并到单个 FFI 调用。

现状

_MotrixSimWriteProgram.executemotrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py:303):

for compiled, sub_buffers in self._ops:
    compiled.op(sub_buffers, idx, rows)
if self._refresh_kinematics:
    self._model.forward_kinematic(rows)
  • N 个 write 声明 = N 次 op 调用,每个 write 一块独立 buffer
  • _MultiTargetOp 对 body/actuator/link/geom/mocap 逐 target 调 Python setter(set_translation / set_rotation / set_global_linear_velocity / set_kp_override / set_mass_override / set_friction_override / set_pose 等),K 个 target 就是 K 次 Python→C FFI 调用
  • forward kinematics 是额外一次单独 FFI 调用
  • 写冲突检测(_claim / _claim_ctrl_targets)在 Python 侧逐索引执行

Proposal

将语义类 write 编译进一个 compile_write plan,执行时一次 program.execute(data, env_ids, forward_kinematic=...) 完成 scatter,FK 由 native 每至多刷一次:

中性声明 motrixsim write.WriteSource
CtrlTargetsWrite ActuatorCtrls
BodyPositionWrite / BodyRotationWrite BodyPosition / BodyRotation(含 mocap)
BodyLinearVelocityWrite / BodyAngularVelocityWrite BodyLinearVelocity / BodyAngularVelocity
ActuatorKpWrite / ActuatorDampingWrite ActuatorKpOverride / ActuatorDampingOverride
BodyMassWrite / BodyComWrite LinkMassOverride / LinkCenterOfMassOverride
GeomFrictionWrite GeomFrictionOverride
  • 无原生 source、保留现有 numpy scatter 的:DofPositionWrite / DofVelocityWrite / BodyJointPositionWrite / BodyJointVelocityWrite / JointPositionWrite / JointVelocityWrite(DOF 通道直写 rows.dof_pos/dof_vel,本无 FFI 开销);MocapPoseWrite 可拆为 BodyPosition + BodyRotation 两个 source 合入同一 plan,或暂保留现实现
  • native 编译期原生拒绝重复写同一 channel,可替代 Python 侧 _claim / _claim_ctrl_targets(注意保留现有错误信息语义,测试有断言)
  • reset 路径 _execute_reset 已是单次 rows.reset(...),可原样保留
  • 中性层 WriteProgram.buffer(name) 接口不变:native WriteProgram 提供连续 buffer 的零拷贝字段视图 program[name],可直接作为 buffer(name) 的返回

收益

执行阶段从「每 target 一次 FFI」降为「每步 1 次 FFI + 若干 numpy 赋值」;FK 合并进同一次调用;冲突检测移入 native 编译期。

References

  • motrix_env_motrixsim/src/motrix_env_motrixsim/write_compiler.py
  • motrix_env_motrixsim/tests/test_write_compiler.pytest_multi_target_writes.py(迁移后需保持通过)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions